From aa00450c63846e625d56373cd3ba0ffa54e22726 Mon Sep 17 00:00:00 2001 From: Shreyas Garimella Date: Fri, 26 Jun 2026 11:36:50 -0700 Subject: [PATCH 01/15] Add EgoDex hand-pose scenario search example Read Apple's EgoDex through Daft's LeRobot v3 reader, precompute per-frame geometric states + Daft-window action rates and SigLIP-2 semantic embeddings, and serve a query UI that combines a geometric filter with semantic ranking. Dataset is CC-BY-NC-ND: scripts contain no EgoDex data; users convert their own licensed copy locally. Co-Authored-By: Claude Opus 4.8 (1M context) --- datasets/egodex/README.md | 78 +++ datasets/egodex/clip_features.py | 84 ++++ datasets/egodex/egodex_lerobot.py | 181 +++++++ datasets/egodex/hdf5.py | 561 ++++++++++++++++++++++ datasets/egodex/lerobot.py | 338 +++++++++++++ datasets/egodex/pose_features.py | 118 +++++ datasets/egodex/query_ui.py | 518 ++++++++++++++++++++ datasets/egodex/requirements-clip.txt | 15 + datasets/egodex/requirements-ui.txt | 1 + datasets/egodex/run_clip_features.py | 72 +++ datasets/egodex/run_pose_features.py | 165 +++++++ datasets/egodex/skeleton_features.py | 212 ++++++++ datasets/egodex/test_skeleton_features.py | 155 ++++++ 13 files changed, 2498 insertions(+) create mode 100644 datasets/egodex/README.md create mode 100644 datasets/egodex/clip_features.py create mode 100644 datasets/egodex/egodex_lerobot.py create mode 100644 datasets/egodex/hdf5.py create mode 100644 datasets/egodex/lerobot.py create mode 100644 datasets/egodex/pose_features.py create mode 100644 datasets/egodex/query_ui.py create mode 100644 datasets/egodex/requirements-clip.txt create mode 100644 datasets/egodex/requirements-ui.txt create mode 100644 datasets/egodex/run_clip_features.py create mode 100644 datasets/egodex/run_pose_features.py create mode 100644 datasets/egodex/skeleton_features.py create mode 100644 datasets/egodex/test_skeleton_features.py diff --git a/datasets/egodex/README.md b/datasets/egodex/README.md new file mode 100644 index 0000000..b162947 --- /dev/null +++ b/datasets/egodex/README.md @@ -0,0 +1,78 @@ +# EgoDex: scenario search over hand-pose video with Daft + +Find specific physical **states** and **actions** in egocentric hand-manipulation video — +"a writing-gripped hand", "the wrist twisting", "fingers closing into a grasp" — that a +text-only search can't reach. This example reads [Apple's EgoDex](https://github.com/apple/ml-egodex) +through Daft's LeRobot v3 reader, precomputes per-frame geometry and SigLIP semantic embeddings, +and serves a query UI that combines a geometric filter with semantic ranking. + +``` +raw EgoDex HDF5 + │ egodex_lerobot.py # convert → LeRobot v3 (48-D state + 204-D skeleton + video) + ▼ +LeRobot v3 dataset + │ run_pose_features.py # NumPy states + Daft-window action rates → out/pose_features + │ run_clip_features.py # SigLIP-2 image embeddings (GPU) → out/clip_features + ▼ +query_ui.py # filter by pose scenario + rank by semantic text +``` + +## ⚠️ Dataset license + +EgoDex is released under **CC-BY-NC-ND**: non-commercial use only, and **no redistribution of +derivatives**. These scripts contain no EgoDex data — you must obtain your own copy from +[apple/ml-egodex](https://github.com/apple/ml-egodex) and run the conversion locally. Do not +redistribute the converted dataset or clips derived from it. + +## Files + +| File | Purpose | +| --- | --- | +| `hdf5.py` | Vendored `daft.datasets` HDF5 reader. | +| `lerobot.py` | Vendored `daft.datasets.lerobot` v3 reader (one row per frame, lazy video decode). | +| `egodex_lerobot.py` | Convert raw EgoDex HDF5 → LeRobot v3, emitting 48-D `observation.state` and 204-D `observation.skeleton`. | +| `pose_features.py` | Per-frame features from the 48-D state (curl, wrist, palm normal, pinch, aperture). | +| `skeleton_features.py` | Per-frame features from the 204-D skeleton (finger flexion, palm plane, arm extension, grip predicates). | +| `run_pose_features.py` | The pipeline: NumPy per-frame **states** + **action** rates via Daft window functions → one parquet. | +| `clip_features.py` | SigLIP-2 image/text embedding UDF (`@daft.cls`). | +| `run_clip_features.py` | Embed frames once with SigLIP → embeddings parquet (the only GPU step). | +| `query_ui.py` | Gradio demo: filter by pose scenario, rank by semantic text, play matched segments. | +| `test_skeleton_features.py` | Verify the 204-D geometry against synthetic + real data. | + +## Run it + +```bash +# 0) Convert your licensed EgoDex HDF5 → a local LeRobot v3 dataset (one-time) +python egodex_lerobot.py # writes ./egodex_lerobot_full + +# 1) Geometry precompute — Daft windows, CPU, a few minutes +DATASET=./egodex_lerobot_full python run_pose_features.py # → out/pose_features + +# 2) Semantic embeddings — SigLIP, needs a GPU +DATASET=./egodex_lerobot_full python run_clip_features.py # → out/clip_features + +# 3) Launch the query UI +DATASET=./egodex_lerobot_full python query_ui.py # prints a local + share URL + +# verify the geometry library +DATA=./egodex_lerobot_full python test_skeleton_features.py +``` + +## How it works + +**States** are per-frame geometry — finger flexion, palm orientation, grip type — computed in +NumPy directly from the sensor data, so a "writing grip" is a writing grip regardless of what the +camera sees. **Actions** are rates of change across frames (grasping = fingers closing, twisting = +forearm rolling), computed with Daft window functions: + +```python +per_episode = Window().partition_by("episode_index").order_by("frame_index") +df = df.with_column( + "curl_rate", + (col("curl").lead(1).over(per_episode) - col("curl")) / SECONDS_PER_FRAME, # grasping signal +) +``` + +Both, plus the SigLIP embeddings, are written to parquet once. At query time a pose predicate is a +boolean column scan (milliseconds, no model) and the text query is one dot product against the +stored embeddings — so the geometry narrows candidates and the semantics rank them, with no GPU. diff --git a/datasets/egodex/clip_features.py b/datasets/egodex/clip_features.py new file mode 100644 index 0000000..8de0e3c --- /dev/null +++ b/datasets/egodex/clip_features.py @@ -0,0 +1,84 @@ +"""SigLIP image-embedding UDF for the EgoDex LeRobot pipeline. + +Runs as plain single-node Daft on a laptop CPU or a GPU box. Encodes each frame +ONCE into a unit-norm SigLIP image embedding (``clip_emb``) and stores it, so any +scenario query later is just a cheap cosine similarity against a text embedding. +""" + +from __future__ import annotations + +import os + +import daft +import numpy as np +import torch +from daft import DataType, Series +from PIL import Image +from transformers import AutoModel, AutoProcessor + +# --- device / config ------------------------------------------------------ +# Resolve the device ONCE. CUDA on the EC2 GPU box; CPU/MPS for a local smoke +# test. gpus=0 when there's no CUDA so @daft.cls doesn't demand a GPU on a CPU +# machine. Force with CLIP_DEVICE=cpu (e.g. to test the CPU path on a Mac). +_HAS_CUDA = torch.cuda.is_available() +GPUS = 1 if _HAS_CUDA else 0 + + +def _auto_device() -> str: + if _HAS_CUDA: + return "cuda" + if torch.backends.mps.is_available(): + return "mps" + return "cpu" + + +DEVICE = os.environ.get("CLIP_DEVICE", _auto_device()) +DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32 + +# Episodes to embed, identified by episode_index (the dataset's native per-frame +# key). Fill in the indices for your scenarios — e.g. one per task type. +EPISODES: list[int] = [0, 1, 4, 6, 7, 9, 13, 15, 26, 29, 31, 32, 34, 36, 39, 40, 41, 50, 54, 57, 58, 60, 62, 66, 76, 82, 83, 84, 88, 89, 93, 100, 104, 109, 113, 114, 118, 120, 123, 125, 129, 132, 135, 136, 137, 144, 145, 146, 147, 149, 154, 156, 158, 159, 161, 163, 165, 166, 168, 172, 173, 179, 181, 182, 185, 187, 189, 191, 192, 197, 203, 204, 214, 215, 220, 227, 233, 236, 237, 246, 250, 254, 255, 257, 260, 261, 268, 269, 271, 272, 278, 282, 286, 288, 290, 294, 299, 302, 303, 308, 312, 319, 321, 323, 328, 329, 330, 333, 334, 337, 341, 342, 344, 346, 349, 350, 354, 356, 367, 371, 372, 373, 374, 375, 379, 382, 383, 388, 400, 401, 404, 406, 407, 410, 412, 413, 415, 417, 419, 420, 428, 429, 431, 444, 445, 446, 454, 459, 460, 464, 468, 470, 471, 472, 474, 475, 479, 480, 484, 487, 490, 493, 494, 497, 504, 505, 507, 509, 511, 512, 520, 521, 529, 531, 533, 534, 540, 541, 542, 543, 546, 550, 552, 564, 567, 570, 571, 572, 575, 576, 577, 581, 582, 583, 584, 586, 596, 599, 601, 604, 606, 607, 609, 612, 618, 619, 620, 628, 630, 632, 637, 638, 639, 640, 643, 645, 648, 655, 658, 663, 667, 669, 670, 671, 678, 682, 689, 693, 698, 710, 715, 716, 728, 730, 735, 738, 745, 746, 753, 757, 759, 764, 765, 769, 770, 772, 774, 775, 778, 779, 780, 782, 783, 786, 787, 789, 797, 798, 802, 803, 810, 811, 813, 816, 826, 827, 832, 833, 834, 835, 836, 837, 838, 842, 844, 848, 851, 852, 855, 858, 861, 862, 874, 881, 882, 883, 884, 889, 892, 894, 895, 896, 901, 903, 904, 905, 908, 909, 912, 913, 921, 928, 940, 944, 951, 953, 954, 956, 958, 964, 969, 977, 978, 980, 981, 983, 991, 994, 995, 999, 1015, 1016, 1020, 1026, 1027, 1034, 1042, 1043, 1046, 1050, 1052, 1057, 1060, 1061, 1063, 1064, 1065, 1066, 1068, 1073, 1075, 1089, 1104, 1106, 1107, 1110, 1111, 1125, 1127, 1128, 1134, 1140, 1141, 1146, 1147, 1150, 1154, 1158, 1160, 1165, 1175, 1178, 1179, 1185, 1186, 1187, 1189, 1192, 1197, 1198, 1201, 1203, 1205, 1207, 1213, 1218, 1220, 1222, 1224, 1225, 1226, 1229, 1230, 1233, 1234, 1236, 1238, 1239, 1240, 1242, 1250, 1251, 1260, 1261, 1270, 1275, 1278, 1289, 1293, 1294, 1295, 1298, 1303, 1306, 1310, 1313, 1318, 1320, 1322, 1326, 1328, 1329, 1331, 1332, 1335, 1339, 1345, 1352, 1354, 1362, 1364, 1365, 1368, 1369, 1372, 1373, 1375, 1380, 1381, 1382, 1383, 1385, 1390, 1392, 1393, 1394, 1402, 1407, 1410, 1414, 1416, 1419, 1421, 1424, 1430, 1432, 1434, 1442, 1443, 1444, 1446, 1449, 1450, 1452, 1455, 1460, 1462, 1463, 1466, 1468, 1471, 1478, 1479, 1481, 1483, 1489, 1498, 1502, 1504, 1508, 1510, 1511, 1512, 1516, 1520, 1521, 1522, 1523, 1525, 1538, 1541, 1542, 1548, 1550, 1552, 1562, 1565, 1570, 1571, 1573, 1576, 1577, 1583, 1587, 1589, 1590, 1591, 1592, 1595, 1601, 1602, 1605, 1613, 1617, 1618, 1622, 1624, 1630, 1631, 1633, 1643, 1645, 1647, 1651, 1652, 1658, 1660, 1662, 1671, 1673, 1678, 1682, 1684, 1685, 1697, 1698, 1699, 1700, 1702, 1707, 1708, 1710, 1711, 1713, 1717, 1719, 1720, 1721, 1722, 1727, 1728, 1729, 1732, 1733, 1737, 1743, 1744, 1745, 1746, 1747, 1750, 1752, 1763, 1765, 1766, 1768, 1777, 1778, 1782, 1783, 1786, 1799, 1807, 1811, 1813, 1818, 1819, 1820, 1821, 1822, 1824, 1827, 1828, 1829, 1832, 1834, 1835, 1837, 1839, 1845, 1849, 1855, 1857, 1858, 1859, 1862, 1865, 1866, 1875, 1877, 1878, 1879, 1887, 1892, 1894, 1909, 1912, 1913, 1916, 1926, 1928, 1930, 1931, 1933, 1936, 1938, 1941, 1948, 1951, 1952, 1953, 1958, 1962, 1970, 1971, 1973, 1975, 1979, 1990, 1991, 1992, 1995, 1998, 2001, 2002, 2004, 2007, 2008, 2014, 2015, 2017, 2018, 2021, 2025, 2029, 2034, 2036, 2038, 2041, 2042, 2054, 2057, 2062, 2067, 2068, 2076, 2077, 2080, 2082, 2083, 2094, 2097, 2114, 2115, 2121, 2125, 2129, 2130, 2133, 2134, 2135, 2137, 2138, 2142, 2143, 2145, 2146, 2148, 2150, 2153, 2154, 2157, 2160, 2164, 2167, 2168, 2170, 2171, 2176, 2181, 2184, 2185, 2188, 2193, 2195, 2203, 2213, 2214, 2223, 2224, 2226, 2227, 2229, 2233, 2238, 2241, 2242, 2246, 2247, 2249, 2255, 2257, 2263, 2265, 2279, 2286, 2287, 2292, 2296, 2297, 2299, 2309, 2312, 2318, 2324, 2327, 2331, 2333, 2335, 2339, 2342, 2350, 2351, 2357, 2359, 2374, 2379, 2382, 2387, 2389, 2395, 2397, 2401, 2404, 2406, 2415, 2417, 2421, 2429, 2435, 2442, 2443, 2444, 2446, 2448, 2450, 2451, 2456, 2458, 2462, 2465, 2466, 2468, 2469, 2470, 2473, 2477, 2480, 2482, 2483, 2488, 2491, 2497, 2499, 2502, 2505, 2506, 2507, 2509, 2516, 2519, 2523, 2528, 2532, 2535, 2541, 2542, 2543, 2547, 2549, 2553, 2559, 2561, 2562, 2566, 2567, 2569, 2570, 2571, 2580, 2581, 2588, 2589, 2593, 2595, 2597, 2600, 2601, 2610, 2614, 2616, 2622, 2623, 2625, 2626, 2629, 2638, 2641, 2648, 2649, 2650, 2658, 2659, 2665, 2677, 2684, 2687, 2689, 2693, 2696, 2701, 2702, 2704, 2705, 2709, 2712, 2715, 2716, 2719, 2726, 2728, 2736, 2739, 2742, 2745, 2748, 2753, 2755, 2758, 2759, 2760, 2762, 2767, 2769, 2778, 2779, 2780, 2782, 2783, 2786, 2787, 2788, 2789, 2790, 2795, 2797, 2799, 2801, 2804, 2806, 2810, 2811, 2813, 2820, 2827, 2829, 2834, 2844, 2846, 2856, 2857, 2858, 2859, 2861, 2862, 2864, 2866, 2868, 2869, 2870, 2873, 2874, 2878, 2882, 2883, 2886, 2888, 2889, 2890, 2894, 2895, 2896, 2897, 2898, 2903, 2904, 2907, 2908, 2909, 2916, 2922, 2925, 2927, 2931, 2932, 2934, 2941, 2943, 2947, 2950, 2955, 2956, 2960, 2969, 2970, 2971, 2977, 2978, 2979, 2980, 2981, 2982, 2989, 2991, 2996, 2999, 3004, 3006, 3007, 3010, 3011, 3026, 3028, 3032, 3034, 3035, 3036, 3038, 3042, 3048, 3049, 3051, 3053, 3054, 3057, 3062, 3063, 3065, 3068, 3074, 3075, 3077, 3081, 3083, 3084, 3085, 3089, 3095, 3096, 3102, 3103, 3104, 3109, 3123, 3125, 3128, 3129, 3131, 3132, 3138, 3140, 3142, 3144, 3149, 3150, 3155, 3159, 3163, 3174, 3177, 3179, 3185, 3190, 3191, 3193, 3196, 3205, 3207, 3208, 3209, 3210, 3211, 3216, 3217, 3218, 3219, 3221, 3222, 3230, 3232] + +SUBSAMPLE = 30 # keep 1 of every 30 frames (~1 fps); semantic content barely changes between adjacent frames +MODEL_ID = "google/siglip2-base-patch16-224" +EMB_DIM = 768 # SigLIP2-base shared image/text embedding dim (must match the model) + + +# --- SigLIP embedding ------------------------------------------------------ +def _normalized_embedding(model_output) -> torch.Tensor: + """Pull the embedding tensor out of a transformers output and L2-normalize it. + + transformers 5.x returns a model-output object from get_image_features / + get_text_features; older versions returned a bare tensor. Handle both. + """ + if torch.is_tensor(model_output): + feats = model_output + else: + feats = model_output.pooler_output + feats = feats.float() + return feats / feats.norm(dim=-1, keepdim=True) # unit-norm so cosine == dot product + + +@daft.cls(gpus=GPUS, max_concurrency=1, use_process=False) +class SiglipEmbedder: + """Encode each frame into a unit-norm SigLIP image embedding. + + Single-node, in-process: runs on this machine's device (CUDA on the EC2 GPU + box, CPU/MPS for a local smoke test). The model loads once per instance. + """ + + def __init__(self) -> None: + self.model = AutoModel.from_pretrained(MODEL_ID, torch_dtype=DTYPE).to(DEVICE).eval() + self.processor = AutoProcessor.from_pretrained(MODEL_ID) + + @daft.method.batch(return_dtype=DataType.embedding(DataType.float32(), EMB_DIM), batch_size=16) + def embed_image(self, images: Series): + pil_frames = [] + for array in images.to_pylist(): + pil_frames.append(Image.fromarray(np.asarray(array, dtype=np.uint8))) + + inputs = self.processor(images=pil_frames, return_tensors="pt").to(DEVICE) + with torch.no_grad(): + model_output = self.model.get_image_features(**inputs) + embeddings = _normalized_embedding(model_output) + return list(embeddings.cpu().numpy()) diff --git a/datasets/egodex/egodex_lerobot.py b/datasets/egodex/egodex_lerobot.py new file mode 100644 index 0000000..50c569e --- /dev/null +++ b/datasets/egodex/egodex_lerobot.py @@ -0,0 +1,181 @@ +# /// script +# description = "Convert raw EgoDex HDF5 episodes into a LeRobot v3 dataset (48-D state + 204-D skeleton)" +# requires-python = ">=3.10, <3.13" +# dependencies = ["daft>=0.7.15", "numpy", "h5py", "lerobot"] +# /// +"""EgoDex (raw HDF5) → LeRobot-format frame features, in Daft. + +Reads raw EgoDex HDF5 with daft.datasets.hdf5.read and produces the per-frame +LeRobot feature columns (observation.state[48], observation.extrinsics[16], +action[48], task), then writes an on-disk LeRobot v3 dataset (tabular only — the +observation.image video feature is added separately by egodex_video.py). + +""" + +from __future__ import annotations +import numpy as np +import daft +from daft import col +from daft.datatype import DataType +from daft.datasets import hdf5 +from daft.functions import coalesce, when +from daft.udf import func +from daft.window import Window +import pathlib +import shutil + +FPS = 30.0 + +# observation.state is 48 floats: per hand, wrist xyz + rot6d (first two rotation +# columns) + 5 fingertip xyz (thumb..little); left hand then right hand. +WRIST = {"left": "transforms/leftHand", "right": "transforms/rightHand"} +TIPS = { + "left": [ + "transforms/leftThumbTip", + "transforms/leftIndexFingerTip", + "transforms/leftMiddleFingerTip", + "transforms/leftRingFingerTip", + "transforms/leftLittleFingerTip", + ], + "right": [ + "transforms/rightThumbTip", + "transforms/rightIndexFingerTip", + "transforms/rightMiddleFingerTip", + "transforms/rightRingFingerTip", + "transforms/rightLittleFingerTip", + ], +} +CAMERA = "transforms/camera" +# The 12 transforms feeding build_state, in the order it expects them. +STATE_TRANSFORMS = [ + WRIST["left"], + *TIPS["left"], + WRIST["right"], + *TIPS["right"], +] +ATTRS = ["llm_description", "llm_description2", "which_llm_description"] + +# observation.skeleton is 204 floats: the xyz translation of every joint, in +# joint order. Per side: hand, forearm, arm, shoulder, then each finger chain +# (thumb has 4 parts; index/middle/ring/little add a metacarpal); then body +# joints (hip, spine1-7, neck1-4). Camera is excluded (it is observation.extrinsics). +FINGERS = ["Thumb", "Index", "Middle", "Ring", "Little"] + +def finger_transforms(side, finger): + infix = "" if finger == "Thumb" else "Finger" + parts = (["Metacarpal"] if finger != "Thumb" else []) + ["Knuckle", "IntermediateBase", "IntermediateTip", "Tip"] + return [f"transforms/{side}{finger}{infix}{part}" for part in parts] + +def side_transforms(side): + arm = [f"transforms/{side}{j}" for j in ("Hand", "Forearm", "Arm", "Shoulder")] + fingers = [t for finger in FINGERS for t in finger_transforms(side, finger)] + return arm + fingers + +BODY_TRANSFORMS = [f"transforms/{j}" for j in ("hip", *(f"spine{i}" for i in range(1, 8)), *(f"neck{i}" for i in range(1, 5)))] +SKELETON_TRANSFORMS = side_transforms("left") + side_transforms("right") + BODY_TRANSFORMS +SKELETON_DIM = len(SKELETON_TRANSFORMS) * 3 + +# Transforms every convertible episode must contain (used by egodex_preflight). +REQUIRED_TRANSFORMS = SKELETON_TRANSFORMS + [CAMERA] + +def hand_block(wrist, tips): + # wrist xyz (translation) + rot6d (first two rotation columns) + each fingertip xyz + w = np.asarray(wrist, dtype=np.float32).reshape(4, 4) + out = list(w[:3, 3]) + list(w[:3, 0]) + list(w[:3, 1]) + for tip in tips: + out += list(np.asarray(tip, dtype=np.float32).reshape(4, 4)[:3, 3]) + return out + +@func(return_dtype=DataType.tensor(DataType.float32(), shape=(16,))) +def build_extrinsics(camera): + return np.asarray(camera, dtype=np.float32).reshape(16) # camera 4x4, row-major + +@func(return_dtype=DataType.tensor(DataType.float32(), shape=(48,))) +def build_state(left_hand, left_thumb, left_index, left_middle, left_ring, left_little, + right_hand, right_thumb, right_index, right_middle, right_ring, right_little): + block = hand_block(left_hand, [left_thumb, left_index, left_middle, left_ring, left_little]) + block += hand_block(right_hand, [right_thumb, right_index, right_middle, right_ring, right_little]) + return np.asarray(block, dtype=np.float32) + +@func(return_dtype=DataType.tensor(DataType.float32(), shape=(SKELETON_DIM,))) +def build_skeleton(*joints): + # Each joint is a 4x4 transform; take its xyz translation, in SKELETON_TRANSFORMS order. + return np.concatenate([np.asarray(j, dtype=np.float32).reshape(4, 4)[:3, 3] for j in joints]) + +def egodex_frames(path): + """Read raw EgoDex HDF5 → per-frame LeRobot feature columns (no video).""" + df = hdf5.read(path, datasets=SKELETON_TRANSFORMS + [CAMERA], attrs=ATTRS) + df = df.with_column("observation.state", build_state(*[col(n) for n in STATE_TRANSFORMS])) + df = df.with_column("observation.skeleton", build_skeleton(*[col(n) for n in SKELETON_TRANSFORMS])) + df = df.with_column("observation.extrinsics", build_extrinsics(col(CAMERA))) + # task = llm_description, or llm_description2 for reversible tasks (which_llm_description="2"); + # fill_null(False) makes the absent-which case fall through to llm_description. + df = df.with_column( + "task", + when((col("which_llm_description") == "2").fill_null(False), col("llm_description2")) + .otherwise(col("llm_description")), + ) + # action = next frame's state within the episode; the last frame repeats itself (no wrap). + window = Window().partition_by("path").order_by("row_index") + df = df.with_column("action", coalesce(col("observation.state").lead(1).over(window), col("observation.state"))) + return df.select("path", "row_index", "observation.state", "observation.skeleton", "observation.extrinsics", "action", "task") + +def state_names(): + # The 48 per-dimension names for observation.state. + names = [] + for side in ("left", "right"): + names += [f"{side}_wrist_{a}" for a in "xyz"] + names += [f"{side}_rot_{i}" for i in range(6)] + for finger in ("thumb", "index", "middle", "ring", "little"): + names += [f"{side}_{finger}_{a}" for a in "xyz"] + return names + +def skeleton_names(): + # The 204 per-dimension names for observation.skeleton (joint xyz, in transform order). + return [f"{t.split('/')[-1]}_{a}" for t in SKELETON_TRANSFORMS for a in "xyz"] + +def features(): + return { + "observation.state": {"dtype": "float32", "shape": (48,), "names": state_names()}, + "observation.skeleton": {"dtype": "float32", "shape": (SKELETON_DIM,), "names": skeleton_names()}, + "observation.extrinsics": {"dtype": "float32", "shape": (16,), "names": [f"extrinsic_{i}" for i in range(16)]}, + "action": {"dtype": "float32", "shape": (48,), "names": [f"action_{i}" for i in range(48)]}, + } + +def write_lerobot(files, repo_id, output_dir, batch_size=64): + """Write EgoDex HDF5 episodes to an on-disk LeRobot v3 dataset (tabular only, no video). + + Daft reads + transforms a batch of files in parallel (one file = one episode); each + episode's frames are then fed in order to the serial LeRobotDataset writer. The batch + size bounds driver memory, and the action window partitions by path so episodes in a + batch never bleed together. + """ + from lerobot.datasets.lerobot_dataset import LeRobotDataset # heavy optional dep; only needed to write + + out = pathlib.Path(output_dir) + if out.exists(): + shutil.rmtree(out) + ds = LeRobotDataset.create( + repo_id=repo_id, fps=int(FPS), features=features(), root=str(out), robot_type="hand", use_videos=False + ) + episodes = 0 + for start in range(0, len(files), batch_size): + rows = egodex_frames(files[start : start + batch_size]).sort(["path", "row_index"]).to_pydict() + paths = rows["path"] + n = len(paths) + i = 0 + while i < n: # group consecutive rows sharing a path into one episode + p = paths[i] + while i < n and paths[i] == p: + ds.add_frame({ + "observation.state": np.asarray(rows["observation.state"][i], dtype=np.float32), + "observation.skeleton": np.asarray(rows["observation.skeleton"][i], dtype=np.float32), + "observation.extrinsics": np.asarray(rows["observation.extrinsics"][i], dtype=np.float32), + "action": np.asarray(rows["action"][i], dtype=np.float32), + "task": rows["task"][i], + }) + i += 1 + ds.save_episode() + episodes += 1 + ds.finalize() + return str(out), episodes diff --git a/datasets/egodex/hdf5.py b/datasets/egodex/hdf5.py new file mode 100644 index 0000000..baa6064 --- /dev/null +++ b/datasets/egodex/hdf5.py @@ -0,0 +1,561 @@ +"""Generic HDF5 readers for `daft.datasets`. + +An HDF5 file is like a filesystem inside a single file: it contains *groups* +(folders), *datasets* (n-dimensional arrays, each with a numpy dtype and a +shape), and *attributes* (small key-value tags attached to groups/datasets). A +very common layout -- used by e.g. EgoDex (https://github.com/apple/ml-egodex) +-- is "one file per episode", where most datasets share a leading time +dimension ``N`` (one entry per frame) plus a few constant datasets that don't: + + camera/intrinsic (3, 3) constant per file + transforms/leftHand (N, 4, 4) one 4x4 pose per frame + confidences/leftHand (N,) one scalar per frame + +and per-file root attributes (e.g. EgoDex's ``llm_description`` task text). + +:func:`read` maps that layout onto a DataFrame with one row per frame: +datasets whose leading dimension matches the file's "row axis" contribute one +element per row, everything else is broadcast (repeated) onto every row, and +requested ``attrs`` are broadcast as additional columns. +:func:`read_datasets` is the exploration view -- one row per (file, dataset) +with its name/shape/dtype -- useful for deciding what to pass to :func:`read`. +""" + +from __future__ import annotations + +import io +from collections import Counter +from typing import TYPE_CHECKING, Any, NamedTuple + +import daft +from daft.api_annotations import PublicAPI +from daft.datatype import DataType +from daft.expressions import col +from daft.file import File +from daft.functions.file_ import file as file_expr +from daft.udf import func + +if TYPE_CHECKING: + import h5py + import numpy as np + + from daft.daft import IOConfig + from daft.dataframe import DataFrame + + +# Files up to this size are read fully into memory before parsing. HDF5 reads +# metadata in many tiny seek+read calls, and each call on a Daft file handle +# crosses the Python<->Rust boundary individually (~14x slower than reading +# from RAM). Above the threshold we stream instead to bound memory usage. +_SLURP_MAX_BYTES = 256 * 1024 * 1024 + +# Schema of one read_datasets entry (per file, per dataset). +_DATASET_ENTRY_DTYPE = DataType.struct( + {"name": DataType.string(), "shape": DataType.list(DataType.int64()), "dtype": DataType.string()} +) + + +class _ColumnSpec(NamedTuple): + """How one HDF5 dataset becomes one DataFrame column. + + per_frame=True -> the dataset's leading dim is the row axis; row i gets + element ``arr[i]``. + per_frame=False -> the dataset is a per-file constant; every row gets the + whole array (broadcast). + """ + + name: str + dtype: DataType + per_frame: bool + is_string: bool + elem_shape: tuple[int, ...] + np_dtype: str + + +class _AttrSpec(NamedTuple): + """How one HDF5 root attribute becomes one (broadcast) DataFrame column. + + Attributes are per-file scalars/strings/small arrays, so each one is + repeated onto every row of its file. ``kind`` is a small tag the loader + maps to a pyarrow type (see ``_arrow_type_for_kind``). + """ + + name: str + dtype: DataType + kind: str + + +def _import_h5py() -> Any: + try: + import h5py + except ImportError as err: + raise ImportError("`daft.datasets.hdf5` requires h5py. Install with `pip install h5py`.") from err + return h5py + + +def _open_h5(f: Any) -> h5py.File: + """Open a file-like object as HDF5, preferring an in-memory copy for small files. + + h5py can read directly from any Python file object that supports + read/seek/tell, but issues many small reads, which are slow against a Daft + file handle. Small files are slurped into a BytesIO instead; large files + are read in place, falling back to a slurp if h5py's driver rejects the + handle. + """ + h5py = _import_h5py() + size = f.seek(0, io.SEEK_END) + f.seek(0) + if size <= _SLURP_MAX_BYTES: + return h5py.File(io.BytesIO(f.read()), "r") + try: + return h5py.File(f, "r") + except (OSError, ValueError, TypeError): + f.seek(0) + return h5py.File(io.BytesIO(f.read()), "r") + + +def _walk_datasets(h5: h5py.File) -> list[tuple[str, tuple[int, ...], np.dtype]]: + """Collect (path, shape, dtype) for every dataset in the file. + + Groups are walked like directories; the leaves (datasets) are recorded by + their full slash-separated path, e.g. ``transforms/leftHand``, in sorted + order. Only metadata is read, never array data. + """ + h5py = _import_h5py() + found: list[tuple[str, tuple[int, ...], Any]] = [] + groups: list[Any] = [h5] + while groups: + group = groups.pop() + for key in group: + obj = group[key] + if isinstance(obj, h5py.Dataset): + # obj.name is the absolute path, e.g. "/transforms/leftHand" + found.append((obj.name.lstrip("/"), tuple(obj.shape), obj.dtype)) + elif isinstance(obj, h5py.Group): + groups.append(obj) + found.sort(key=lambda entry: entry[0]) + return found + + +def _element_datatype(np_dtype: Any, elem_shape: tuple[int, ...]) -> DataType | None: + """Map one dataset element (one row's worth of data) to a Daft DataType. + + Returns None for dtypes we can't represent (e.g. compound/struct dtypes), + so callers can skip those datasets instead of failing. + """ + h5py = _import_h5py() + if h5py.check_string_dtype(np_dtype): + # Variable-length strings only make sense as scalars-per-row here. + return DataType.string() if elem_shape == () else None + try: + scalar = DataType.from_numpy_dtype(np_dtype) + except Exception: + return None + if elem_shape == (): + return scalar + return DataType.tensor(scalar, shape=elem_shape) + + +def _normalize_attr(raw: Any) -> Any: + """Convert a raw h5py attribute value to a plain Python value. + + h5py returns attributes as numpy scalars/arrays or bytes; we normalize to + str / bool / int / float / list so they can go straight into a pyarrow + column. Returns ``str(raw)`` for anything exotic. + """ + import numpy as np + + if raw is None: + return None + if isinstance(raw, np.ndarray): + if raw.ndim == 0: + raw = raw.item() + else: + return [_normalize_attr(x) for x in raw.tolist()] + if isinstance(raw, np.generic): + raw = raw.item() + if isinstance(raw, bytes): + return raw.decode("utf-8", "replace") + if isinstance(raw, (bool, int, float, str)): + return raw + return str(raw) + + +def _attr_kind_dtype(value: Any) -> tuple[str, DataType]: + """Infer the (kind tag, Daft DataType) for a normalized attribute value.""" + # bool first: in Python `bool` is a subclass of `int`. + if isinstance(value, bool): + return "bool", DataType.bool() + if isinstance(value, int): + return "int", DataType.int64() + if isinstance(value, float): + return "float", DataType.float64() + if isinstance(value, str): + return "string", DataType.string() + if isinstance(value, list): + elem = next((x for x in value if x is not None), None) + if isinstance(elem, bool): + return "list_bool", DataType.list(DataType.bool()) + if isinstance(elem, int): + return "list_int", DataType.list(DataType.int64()) + if isinstance(elem, float): + return "list_float", DataType.list(DataType.float64()) + return "list_string", DataType.list(DataType.string()) + return "string", DataType.string() + + +def _arrow_type_for_kind(kind: str) -> Any: + """Map an attribute ``kind`` tag to a pyarrow type (used inside the loader).""" + import pyarrow as pa + + return { + "bool": pa.bool_(), + "int": pa.int64(), + "float": pa.float64(), + "string": pa.string(), + "list_bool": pa.list_(pa.bool_()), + "list_int": pa.list_(pa.int64()), + "list_float": pa.list_(pa.float64()), + "list_string": pa.list_(pa.string()), + }[kind] + + +def _is_per_frame(name: str, shapes_per_file: list[dict[str, tuple[int, ...]]], row_dims: list[int]) -> bool: + """A dataset is per-frame only if its leading dim matches the row axis in EVERY sampled file. + + This is what disambiguates a constant (3, 3) matrix from 3 frames of + 3-vectors: across two files with different N, only genuinely per-frame + datasets keep tracking N. Files where the dataset is absent (HDF5 groups + like EgoDex's confidences are optional) don't count against it. + """ + for shapes, dim in zip(shapes_per_file, row_dims): + shape = shapes.get(name) + if shape is None: + continue + if len(shape) == 0 or shape[0] != dim: + return False + return True + + +def _load_files_batch( + files: Any, specs: list[_ColumnSpec], attr_specs: list[_AttrSpec], index_column: str | None +) -> Any: + """Read a batch of HDF5 files into one struct column, built directly in pyarrow. + + Each file's row count N comes from the first per-frame dataset present in + it. The struct has one list field per dataset and per requested attribute: + per-frame datasets contribute their N elements, broadcast datasets and + attributes are repeated N times, datasets/attributes missing from a file + are null-filled for its N rows, and ``index_column`` (if given) is a + 0..N-1 counter. + + Building the column with pyarrow instead of returning Python values is a + deliberate performance choice: per-frame tensors would otherwise be + converted to Arrow one tiny numpy object at a time (~70us each, dominating + the read), whereas a flat float buffer wrapped as FixedSizeList is a bulk + copy, and Daft's cast from FixedSizeList to FixedShapeTensor is zero-copy. + """ + import numpy as np + import pyarrow as pa + + h5_files = files.to_pylist() + counts: list[int] = [] + per_file: dict[str, list[Any]] = {spec.name: [] for spec in specs} + per_file_attr: dict[str, list[Any]] = {a.name: [] for a in attr_specs} + for file in h5_files: + with file.open() as f: + with _open_h5(f) as h5: + n = None + for spec in specs: + if spec.per_frame and spec.name in h5: + n = h5[spec.name].shape[0] + break + if n is None: + raise ValueError(f"No per-row datasets present in {file}; cannot determine its row count.") + counts.append(n) + for spec in specs: + if spec.name not in h5: + per_file[spec.name].append(None) + continue + ds = h5[spec.name] + if spec.per_frame and ds.shape[0] != n: + raise ValueError( + f"HDF5 dataset {spec.name!r} was inferred to be per-row, but in {file} it has " + f"leading dim {ds.shape[0]} while the file's row axis is {n}. " + f"If it is a per-file constant, pass broadcast=[{spec.name!r}]." + ) + per_file[spec.name].append(ds.asstr()[...] if spec.is_string else ds[...]) + for attr in attr_specs: + per_file_attr[attr.name].append( + _normalize_attr(h5.attrs[attr.name]) if attr.name in h5.attrs else None + ) + + offsets = pa.array(np.concatenate(([0], np.cumsum(counts, dtype=np.int64))), type=pa.int32()) + names: list[str] = [] + fields: list[Any] = [] + for spec in specs: + chunks = per_file[spec.name] + if spec.is_string: + flat_strings: list[str | None] = [] + for n, chunk in zip(counts, chunks): + if chunk is None: + flat_strings.extend([None] * n) + elif spec.per_frame: + flat_strings.extend(chunk.tolist()) + else: + flat_strings.extend([str(chunk)] * n) + elements = pa.array(flat_strings, type=pa.string()) + else: + elem_size = 1 + for dim in spec.elem_shape: + elem_size *= dim + value_type = pa.from_numpy_dtype(np.dtype(spec.np_dtype)) + elem_type = value_type if spec.elem_shape == () else pa.list_(value_type, elem_size) + parts: list[Any] = [] + for n, chunk in zip(counts, chunks): + if chunk is None: + parts.append(pa.nulls(n, elem_type)) + continue + if spec.per_frame: + flat = chunk.reshape(n, -1) + else: + flat = np.broadcast_to(chunk.reshape(1, -1), (n, elem_size)) + values = pa.array(np.ascontiguousarray(flat).reshape(-1)) + if values.type != value_type: + values = values.cast(value_type) + parts.append(values if spec.elem_shape == () else pa.FixedSizeListArray.from_arrays(values, elem_size)) + elements = pa.concat_arrays(parts) if parts else pa.nulls(0, elem_type) + names.append(spec.name) + fields.append(pa.ListArray.from_arrays(offsets, elements)) + + for attr in attr_specs: + arrow_type = _arrow_type_for_kind(attr.kind) + flat_attr: list[Any] = [] + for n, value in zip(counts, per_file_attr[attr.name]): + flat_attr.extend([value] * n) + elements = pa.array(flat_attr, type=arrow_type) + names.append(attr.name) + fields.append(pa.ListArray.from_arrays(offsets, elements)) + + if index_column is not None: + frame_indices = [np.arange(n, dtype=np.uint64) for n in counts] + flat_index = np.concatenate(frame_indices) if frame_indices else np.empty(0, dtype=np.uint64) + names.append(index_column) + fields.append(pa.ListArray.from_arrays(offsets, pa.array(flat_index))) + + return pa.StructArray.from_arrays(fields, names) + + +@func(return_dtype=DataType.list(_DATASET_ENTRY_DTYPE)) +def _list_file_datasets(file: File) -> list[dict[str, Any]]: + """List the (name, shape, dtype) of every dataset in one HDF5 file.""" + with file.open() as f: + with _open_h5(f) as h5: + return [ + {"name": name, "shape": list(shape), "dtype": str(np_dtype)} + for name, shape, np_dtype in _walk_datasets(h5) + ] + + +def _glob_paths(path: str | list[str], io_config: IOConfig | None) -> DataFrame: + return daft.from_glob_path(path, io_config=io_config).select("path") + + +def _head_paths(paths_df: DataFrame, path: str | list[str], n: int) -> list[str]: + from daft.exceptions import DaftCoreException + + try: + head = paths_df.limit(n).to_pydict()["path"] + except DaftCoreException: + # Daft's glob scan errors out when nothing matches; translate it into + # the error a reader caller would expect. + head = [] + if not head: + raise FileNotFoundError(f"No files matched {path!r}") + return head + + +@PublicAPI +def read( + path: str | list[str], + datasets: list[str] | None = None, + broadcast: list[str] | None = None, + attrs: list[str] | None = None, + io_config: IOConfig | None = None, + index_column: str | None = "row_index", +) -> DataFrame: + """Read HDF5 files as a lazy DataFrame with one row per leading-dimension entry. + + The "row axis" is the most common leading dimension among a file's + datasets (for episode-style data like EgoDex this is the number of frames + N). Datasets that match it contribute one element per row; datasets that + don't (e.g. a constant 3x3 ``camera/intrinsic``) are broadcast onto every + row as whole arrays. + + The schema is inferred by sampling up to two matched files. Two, because a + single file can be ambiguous: a constant ``(3, 3)`` matrix in a 3-frame + file looks identical to per-frame 3-vectors. A second file with a + different frame count resolves this -- per-frame datasets track N, constants + don't. If every sampled file has the same N the ambiguity remains; use + ``broadcast`` to pin such datasets explicitly. + + Args: + path: A file path or glob, or a list of them (e.g. + ``"part1/**/*.hdf5"``). Anything `daft.from_glob_path` accepts. + datasets: Optional subset of dataset paths to read (e.g. + ``["transforms/leftHand", "camera/intrinsic"]``). Defaults to every + dataset with a representable dtype; datasets with unsupported + dtypes (e.g. compound types) are skipped. + broadcast: Dataset paths to always treat as per-file constants + (repeated onto every row), overriding the row-axis inference. + attrs: Names of HDF5 root attributes to surface as broadcast columns + (e.g. ``["llm_description"]``). Each is repeated onto every row of + its file; files missing the attribute get null. Strings, numeric + scalars, booleans, and 1-D arrays of those are supported. + io_config: Optional IO configuration for remote reads. + index_column: Name for the per-file row counter column (the position + along the row axis, i.e. the frame index). Pass None to omit it. + + Returns: + Lazy DataFrame with columns ``path``, the index column, one column per + dataset (named by its slash-separated HDF5 path), and one column per + requested attribute. + """ + paths_df = _glob_paths(path, io_config) + sample_paths = _head_paths(paths_df, path, n=2) + first = sample_paths[0] + + # Schema inference: read dataset metadata + root attributes (no array data) + # from up to two files. "rb" = binary mode; HDF5 is a binary format, the + # text-mode default would try to decode it as UTF-8 and fail. + sampled: list[list[tuple[str, tuple[int, ...], Any]]] = [] + sampled_attrs: list[dict[str, Any]] = [] + for sample_path in sample_paths: + with daft.open_file(sample_path, mode="rb", io_config=io_config) as f: + with _open_h5(f) as h5: + sampled.append(_walk_datasets(h5)) + sampled_attrs.append({k: _normalize_attr(v) for k, v in h5.attrs.items()}) + + # Union the datasets across the sampled files (first occurrence wins for + # shape/dtype): optional groups like EgoDex's confidences may be absent + # from whichever file the glob happens to list first. + seen: set[str] = set() + found: list[tuple[str, tuple[int, ...], Any]] = [] + for file_found in sampled: + for entry in file_found: + if entry[0] not in seen: + seen.add(entry[0]) + found.append(entry) + found.sort(key=lambda entry: entry[0]) + if datasets is not None: + missing = sorted(set(datasets) - {name for name, _, _ in found}) + if missing: + raise ValueError(f"Datasets not found in {first!r}: {missing}") + found = [entry for entry in found if entry[0] in set(datasets)] + if not found: + raise ValueError(f"No datasets to read in {first!r}") + selected_names = {name for name, _, _ in found} + + forced_broadcast = set(broadcast or ()) + unknown_broadcast = sorted(forced_broadcast - selected_names) + if unknown_broadcast: + raise ValueError(f"`broadcast` names not found in {first!r}: {unknown_broadcast}") + + # Resolve requested attributes against the sampled files (first file that + # has the attribute wins for type inference). + attr_specs: list[_AttrSpec] = [] + for name in attrs or (): + present = [sa[name] for sa in sampled_attrs if name in sa and sa[name] is not None] + if present: + kind, dtype = _attr_kind_dtype(present[0]) + else: + # Absent (or null) in the sampled files; it may still appear in + # others (datasets can be heterogeneous, e.g. EgoDex reset tasks + # lack `which_llm_description`). Default to a nullable string column + # rather than failing the read. + kind, dtype = "string", DataType.string() + attr_specs.append(_AttrSpec(name, dtype, kind)) + attr_names = {a.name for a in attr_specs} + if attr_names & selected_names: + raise ValueError(f"`attrs` names collide with dataset names: {sorted(attr_names & selected_names)}") + + # The row axis of each sampled file: the most common leading dimension + # among its (selected) datasets. For EgoDex-style files this is N (frames); + # the lone (3, 3) intrinsic loses the vote. + shapes_per_file: list[dict[str, tuple[int, ...]]] = [ + {name: shape for name, shape, _ in file_found if name in selected_names} for file_found in sampled + ] + row_dims: list[int] = [] + for shapes in shapes_per_file: + votes = Counter(shape[0] for shape in shapes.values() if len(shape) > 0) + if not votes: + raise ValueError(f"All requested datasets in {first!r} are scalars; nothing to use as a row axis.") + row_dims.append(votes.most_common(1)[0][0]) + + specs: list[_ColumnSpec] = [] + for name, shape, np_dtype in found: + per_frame = name not in forced_broadcast and _is_per_frame(name, shapes_per_file, row_dims) + elem_shape = shape[1:] if per_frame else shape + dtype = _element_datatype(np_dtype, elem_shape) + if dtype is None: + continue + specs.append( + _ColumnSpec(name, dtype, per_frame, dtype == DataType.string(), elem_shape, str(np_dtype)) + ) + if not any(spec.per_frame for spec in specs): + raise ValueError(f"No datasets in {first!r} share the inferred row axis (leading dim {row_dims[0]}).") + + # A UDF returns a single column, so pack one list field per dataset/attribute + # into a struct, then unnest + explode so each frame becomes its own row. The + # row counter rides along as its own list field. Field order must match the loader. + struct_fields = {spec.name: DataType.list(spec.dtype) for spec in specs} + for a in attr_specs: + struct_fields[a.name] = DataType.list(a.dtype) + explode_columns = [spec.name for spec in specs] + [a.name for a in attr_specs] + if index_column is not None: + if index_column in selected_names or index_column in attr_names: + raise ValueError(f"index_column {index_column!r} collides with a dataset or attribute name.") + struct_fields[index_column] = DataType.list(DataType.uint64()) + explode_columns.append(index_column) + struct_dtype = DataType.struct(struct_fields) + load_files = func.batch(return_dtype=struct_dtype)(_load_files_batch) + + # Small batches of file paths so the loader UDF runs in parallel across + # files and downstream operators receive rows as soon as the first few + # files are read, instead of waiting for the entire file list. + df = paths_df.into_batches(16) + df = df.with_column( + "_data", load_files(file_expr(col("path"), io_config=io_config), specs, attr_specs, index_column) + ) + df = df.select(col("path"), col("_data").unnest()) + df = df.explode(*explode_columns) + return df + + +@PublicAPI +def read_datasets(path: str | list[str], io_config: IOConfig | None = None) -> DataFrame: + """List every dataset in the matched HDF5 files, one row per (file, dataset). + + This is the exploration view: use it to see what's inside the files + (names, shapes, dtypes) before choosing what to pass to :func:`read`. + + Args: + path: A file path or glob, or a list of them. Anything + `daft.from_glob_path` accepts. + io_config: Optional IO configuration for remote reads. + + Returns: + Lazy DataFrame with columns ``path``, ``name``, ``shape``, ``dtype``. + """ + df = _glob_paths(path, io_config) + df = df.into_batches(16) + df = df.with_column("dataset", _list_file_datasets(file_expr(col("path"), io_config=io_config))) + df = df.explode("dataset") + df = df.select(col("path"), col("dataset").unnest()) + return df + + +__all__ = [ + "read", + "read_datasets", +] diff --git a/datasets/egodex/lerobot.py b/datasets/egodex/lerobot.py new file mode 100644 index 0000000..4f3c9f9 --- /dev/null +++ b/datasets/egodex/lerobot.py @@ -0,0 +1,338 @@ +"""LeRobot Dataset v3.0 helpers for `daft.datasets`. + +This module reads the file-based LeRobot v3 layout (`meta/episodes`, `data`, +`videos`) and exposes episode-level scans plus frame expansion utilities. + +See https://huggingface.co/docs/lerobot/lerobot-dataset-v3 for format details. +""" + +from __future__ import annotations + +import json +import re +from typing import TYPE_CHECKING, Any, TypedDict, cast + +import daft +from daft.api_annotations import PublicAPI +from daft.datatype import DataType +from daft.exceptions import DaftCoreException +from daft.expressions import col, lit +from daft.file import VideoFile +from daft.functions import lpad +from daft.functions.file_ import video_file +from daft.udf import func + +if TYPE_CHECKING: + from daft.daft import IOConfig + from daft.dataframe import DataFrame + + +def _normalize_dataset_root(uri: str) -> str: + """Return a canonical dataset root prefix (no trailing slash) for path joins.""" + u = uri.strip() + # Input looks like a Hugging Face repo ID, i.e. "org/name". + # Paths starting with ".", "/", or a known URI scheme are never repo IDs. + is_hf_repo_id = ( + not u.startswith((".", "/")) + and "://" not in u + and bool(re.fullmatch(r"[\w.-]+/[\w.-]+", uri)) + ) + + if is_hf_repo_id: + u = f"hf://datasets/{u}" + return u.rstrip("/") + + +@func(return_dtype=DataType.image()) +def _decode_lerobot_video_timestamp( + file: VideoFile, + episode_from_timestamp_s: float, + frame_timestamp_s: float, + tolerance_s: float, + image_width_i: int, + image_height_i: int, +) -> Any: # returns a PIL.Image; PIL is an optional dependency imported lazily below + """Pick the decoded frame closest in time to ``episode_from_timestamp_s + frame_timestamp_s``.""" + try: + import av as av_mod + except ImportError as err: + raise ImportError("Decoding LeRobot MP4 shards requires PyAV. Install with `pip install av`.") from err + try: + from PIL import Image as PILImage + except ImportError as err: + raise ImportError( + "Decoding LeRobot MP4 shards requires Pillow. Install with `pip install daft[video]` or `pip install pillow`." + ) from err + abs_ts = float(episode_from_timestamp_s) + float(frame_timestamp_s) + tolerance = float(tolerance_s) + width_i = int(image_width_i) + height_i = int(image_height_i) + width: int | None + height: int | None + if width_i > 0 and height_i > 0: + width, height = width_i, height_i + else: + width, height = None, None + + loaded: list[tuple[float, Any]] = [] + decode_cap = 20_000 + decoded = 0 + + with file.open() as f_open: + with av_mod.open(f_open) as container: + stream = container.streams.video[0] + # Match LeRobot: seek backwards to preceding keyframe, then decode forwards. + container.seek(max(0, int(abs_ts * av_mod.time_base)), backward=True) + + tail_s = max(0.1, tolerance * 50.0, 1.0 / 24.0) + for vf in container.decode(stream): + if vf.pts is None: + continue + current_ts = float(vf.pts * stream.time_base) + pil_img = PILImage.fromarray(vf.to_ndarray(format="rgb24"), mode="RGB") + if width is not None and height is not None: + pil_img = pil_img.resize((width, height), PILImage.Resampling.NEAREST) + + loaded.append((current_ts, pil_img)) + decoded += 1 + + if decoded >= decode_cap: + raise ValueError("Exceeded decode frame budget while aligning to parquet timestamps.") + if current_ts >= abs_ts + tail_s: + break + + if not loaded: + raise ValueError(f"No frames decoded from shard while seeking timestamp {abs_ts:.6f}s.") + + closest_ts, closest_img = min(loaded, key=lambda item: abs(item[0] - abs_ts)) + closest_dist = abs(closest_ts - abs_ts) + if closest_dist > tolerance: + raise ValueError( + f"No frame matched timestamp {abs_ts:.6f}s within tolerance {tolerance} " + f"(closest distance observed: {closest_dist})." + ) + return closest_img + + +class Feature(TypedDict): + dtype: str + + +class LeRobotInfo(TypedDict): + codebase_version: str + data_path: str + video_path: str + fps: float + features: dict[str, Feature] + + +def _read_info(normalized_uri: str, io_config: IOConfig | None = None) -> LeRobotInfo: + with daft.open_file(f"{normalized_uri}/meta/info.json", io_config=io_config) as f: + info = cast("LeRobotInfo", json.load(f)) + if info["codebase_version"] != "v3.0": + raise ValueError("`daft.datasets.lerobot` currently only supports LeRobot datasets of v3 and above") + return info + + +@PublicAPI +def read( + dataset_uri: str, + io_config: IOConfig | None = None, + include_stats: bool = False, + load_video_frames: str | list[str] | bool = False, +) -> DataFrame: + """Read a LeRobot v3 dataset as a lazy DataFrame with one row per frame. + + Reads the per-episode metadata under ``meta/episodes`` and the per-frame + sensor data under ``data``, joins them on ``episode_index``, and broadcasts + each episode's metadata across its frames. Optionally decodes the matching + video frame for one or more camera keys into an image column. + + Args: + dataset_uri: Huggingface repo id (``org/name``), or a local / remote + directory (``s3://...``, ``hf://datasets/...``). + io_config: Optional IO configuration for remote reads. + include_stats: If True, keep the per-episode ``stats/*`` columns + (per-feature min/max/mean/std/quantiles). Defaults to False. + load_video_frames: Which camera keys to decode into image columns, + aligned to each frame's timestamp. Defaults to False (decode + nothing). Pass True to decode every video feature, a single key + (``"observation.image"``), or a list of keys. Decoding requires the + optional ``av`` (PyAV) and ``Pillow`` dependencies. + + Returns: + Lazy DataFrame with one row per frame: the frame's sensor columns, the + broadcast episode metadata, and one image column per decoded video key. + """ + root = _normalize_dataset_root(dataset_uri) + info = _read_info(root, io_config=io_config) + + # Keep the per-episode video metadata (notably `videos/{key}/from_timestamp`, + # the time within the shard where each episode's footage begins). We need it + # to translate episode-local frame timestamps into absolute shard timestamps + # when decoding, and drop these internal columns again before returning. + episode_df = read_episodes( + dataset_uri, io_config=io_config, include_stats=include_stats, include_video_metadata=True + ) + df = load_episode_frames(episode_df, dataset_uri, io_config=io_config) + + # Load video frames into memory + if load_video_frames is not False: + if load_video_frames is True: + video_keys = [name for name, feat_info in info["features"].items() if feat_info["dtype"] == "video"] + elif isinstance(load_video_frames, str): + video_keys = [load_video_frames] + elif isinstance(load_video_frames, list) and all(isinstance(k, str) for k in load_video_frames): + video_keys = load_video_frames + else: + raise ValueError(f"Invalid value provided for argument load_video_frames=`{load_video_frames}`") + + # An MP4 shard packs many episodes back to back, so the shard's internal + # frame numbering is NOT the parquet's episode-local `frame_index` (which + # resets to 0 each episode). Seeking by `frame_index` only happens to work + # for the first episode in each shard. Instead, seek by absolute timestamp: + # `from_timestamp` (where this episode begins in the shard) + the per-frame + # episode-local `timestamp`. That keeps a single coordinate system end to end. + fps = float(info["fps"]) + tolerance_s = 1.0 / fps / 2.0 # half a frame period: any closer frame is unambiguously "the" frame + + # To increase parallelism, reduce batch size + df = df.into_batches(16) + for k in video_keys: + df = df.with_column( + k, + _decode_lerobot_video_timestamp( + col(f"videos/{k}/video"), + col(f"videos/{k}/from_timestamp"), + col("timestamp"), + lit(tolerance_s), + lit(0), # image_width: 0 disables resize (decode at native resolution) + lit(0), # image_height: 0 disables resize + ), + ) + df = df.exclude(f"videos/{k}/video") + + + # Drop the internal per-episode video metadata we kept above (chunk/file index, + # from/to timestamp). This restores read_episodes' default of hiding these. + df = df.exclude(*(c for c in df.column_names if c.startswith("videos/") and not c.endswith("/video"))) + + return df + + +@PublicAPI +def read_episodes( + dataset_uri: str, + io_config: IOConfig | None = None, + include_meta: bool = False, + include_stats: bool = False, + include_video_metadata: bool = False, +) -> DataFrame: + """Read LeRobot v3 episode metadata as a lazy DataFrame (one row per episode). + + This reads the `meta/episodes/**/*.parquet` path under the dataset root. + + Args: + dataset_uri: Huggingface repo id (`org/name`), + or a local / remote directory (`s3://...`, `hf://datasets/...`) + io_config: Optional IO configuration for remote reads. + include_meta: If True, keep the internal ``meta/episodes/*`` columns + (the chunk/file indices locating each episode's own metadata shard). + These are bookkeeping for random access into the sharded metadata + and carry no analytical value once the rows are loaded. Defaults to + False. + include_stats: If True, keep the per-episode ``stats/*`` columns + (per-feature min/max/mean/std/quantiles). Defaults to False. + include_video_metadata: If True, keep the per-episode ``videos/{key}/*`` + columns (the chunk/file indices and from/to timestamps locating each + episode's footage within its video shard). Defaults to False. + + Returns: + Lazy DataFrame of episode metadata, one row per episode. Always includes + a ``videos/{key}/video`` file-handle column per video feature; the + ``include_*`` flags control which additional column families are kept. + """ + root = _normalize_dataset_root(dataset_uri) + info = _read_info(root, io_config=io_config) + df = daft.read_parquet(f"{root}/meta/episodes/**/*.parquet", io_config=io_config) + if not include_meta: + df = df.exclude(*(c for c in df.column_names if c.startswith("meta/"))) + if not include_stats: + df = df.exclude(*(c for c in df.column_names if c.startswith("stats/"))) + + # Get the video keys + video_keys = set(name for name, feat_info in info["features"].items() if feat_info["dtype"] == "video") + + for key in video_keys: + file_name_expr = ( + lit(f"{root}/videos/{key}/chunk-") + + lpad(col(f"videos/{key}/chunk_index").cast(DataType.string), 3, "0") + + lit("/file-") + + lpad(col(f"videos/{key}/file_index").cast(DataType.string), 3, "0") + + lit(".mp4") + ) + + df = df.with_column(f"videos/{key}/video", video_file(file_name_expr, verify=False, io_config=io_config)) + + if not include_video_metadata: + df = df.exclude(*(c for c in df.column_names if c.startswith("videos/") and not c.endswith("/video"))) + + return df + + +@PublicAPI +def load_episode_frames( + episodes: DataFrame, + dataset_uri: str, + io_config: IOConfig | None = None, +) -> DataFrame: + """Expand an episode-level DataFrame into a frame-level DataFrame. + + Reads the per-frame parquet under ``data/**`` and joins it to the provided + episode metadata on ``episode_index``, producing one row per frame. Episode + metadata is broadcast across each episode's frames. + + Filter ``episodes`` before calling this to expand only the episodes you need; + only the surviving episodes contribute to the join. + + Args: + episodes: Episode-level DataFrame, typically from :func:`read_episodes` + (optionally filtered). Must contain an ``episode_index`` column. + dataset_uri: The same dataset identifier passed to :func:`read_episodes` + (Huggingface repo id ``org/name``, or a local / remote directory such + as ``s3://...`` or ``hf://datasets/...``). + io_config: Optional IO configuration for remote reads. + + Returns: + Lazy DataFrame with one row per frame. + """ + root = _normalize_dataset_root(dataset_uri) + + frame_df = daft.read_parquet(f"{root}/data/**/*.parquet", io_config=io_config) + df = episodes.join(frame_df, on=["episode_index"]) + df = df.exclude("data/chunk_index", "data/file_index") + return df + + +@PublicAPI +def read_tasks(dataset_uri: str, io_config: IOConfig | None = None) -> DataFrame: + """Load task metadata as a DataFrame. + + Prefers ``meta/tasks.parquet`` (current LeRobot default). Falls back to legacy + ``meta/tasks.jsonl`` when the Parquet file is missing. + """ + root = _normalize_dataset_root(dataset_uri) + + pq_url = f"{root}/meta/tasks.parquet" + try: + return daft.read_parquet(pq_url, io_config=io_config) + except (OSError, DaftCoreException, FileNotFoundError): + return daft.read_json(f"{root}/meta/tasks.jsonl", io_config=io_config) + + +__all__ = [ + "load_episode_frames", + "read", + "read_episodes", + "read_tasks", +] diff --git a/datasets/egodex/pose_features.py b/datasets/egodex/pose_features.py new file mode 100644 index 0000000..2126091 --- /dev/null +++ b/datasets/egodex/pose_features.py @@ -0,0 +1,118 @@ +"""Per-frame hand-pose features from the 48-D EgoDex observation.state. + +Layout per hand (24 dims): wrist xyz [0:3], wrist rot6d [3:9], 5 fingertips +xyz [9:24] (order [thumb, index, middle, ring, pinky]). Left hand is [0:24], +right hand [24:48]. World +y is up (verified: wrists sit ~0.93 up, left hand -x / +right hand +x, units ~meters). Everything is vectorized over N frames and model-free. +""" +from __future__ import annotations + +import numpy as np + +FINGERS = ["thumb", "index", "middle", "ring", "pinky"] + +# 48-D observation.state layout: per hand, 24 dims = wrist xyz (3) + wrist rot6d (6) +# + 5 fingertips xyz (15). Left hand occupies [0:24], right hand [24:48]. +HAND_BLOCK_DIM = 24 +ROT6D_OFFSET = 3 # rot6d sits at [3:9] within a hand block +ROT6D_LEN = 6 +FINGERTIP_OFFSET = 9 # 5 fingertips xyz sit at [9:24] +HAND_BASE = {"left": 0, "right": HAND_BLOCK_DIM} +HAND_TAGS = [("left", "L"), ("right", "R")] + + +def rot6d_slice(side): + """Column slice of a hand's wrist rot6d (6 values) within the 48-D state ('left' / 'right').""" + start = HAND_BASE[side] + ROT6D_OFFSET + return slice(start, start + ROT6D_LEN) + + +def _split_hand(state, base): + wrist = state[:, base:base + 3] + rot6d = state[:, base + ROT6D_OFFSET:base + ROT6D_OFFSET + ROT6D_LEN] + fingertips = state[:, base + FINGERTIP_OFFSET:base + HAND_BLOCK_DIM].reshape(-1, 5, 3) + return wrist, rot6d, fingertips + + +def palm_normal_from_rot6d(rot6d): + """rot6d = the first two columns of the rotation matrix; the palm normal is their cross product.""" + first_column = rot6d[:, 0:3] + second_column = rot6d[:, 3:6] + first_column = first_column / (np.linalg.norm(first_column, axis=1, keepdims=True) + 1e-9) + second_column = second_column - (first_column * second_column).sum(1, keepdims=True) * first_column + second_column = second_column / (np.linalg.norm(second_column, axis=1, keepdims=True) + 1e-9) + return np.cross(first_column, second_column) + + +def rotation_from_rot6d(rot6d): + """rot6d (N, 6) -> (N, 3, 3) rotation matrices (columns = hand x, y axes + palm normal).""" + first_column = rot6d[:, 0:3] + first_column = first_column / (np.linalg.norm(first_column, axis=1, keepdims=True) + 1e-9) + second_column = rot6d[:, 3:6] + second_column = second_column - (first_column * second_column).sum(1, keepdims=True) * first_column + second_column = second_column / (np.linalg.norm(second_column, axis=1, keepdims=True) + 1e-9) + palm_normal = np.cross(first_column, second_column) + return np.stack([first_column, second_column, palm_normal], axis=2) + + +def compute_raw_features(state: np.ndarray) -> dict: + """Per-frame raw features for both hands, keyed by feature + hand tag ('L' / 'R').""" + features = {} + for side, tag in HAND_TAGS: + wrist, rot6d, fingertips = _split_hand(state, HAND_BASE[side]) + tip_to_wrist = np.linalg.norm(fingertips - wrist[:, None, :], axis=2) # (N, 5) + features[f"fingerdist_{tag}"] = tip_to_wrist + features[f"curl_{tag}"] = tip_to_wrist.mean(1) # small = curled + palm_normal = palm_normal_from_rot6d(rot6d) + features[f"palmnormal_{tag}"] = palm_normal + features[f"palm_up_{tag}"] = palm_normal[:, 1] # +y component + features[f"pinch_{tag}"] = np.linalg.norm(fingertips[:, 0] - fingertips[:, 1], axis=1) # thumb-index + tip_pairs = fingertips[:, :, None, :] - fingertips[:, None, :, :] + features[f"aperture_{tag}"] = np.linalg.norm(tip_pairs, axis=3).max((1, 2)) # max tip-tip spread + features[f"wrist_{tag}"] = wrist + return features + + +# ---- per-episode temporal features (used by the UI's live overlay table) ---- +# The batch precompute computes these rates with Daft window functions instead; +# see run_pose_features.py. These NumPy versions exist only so the UI can fill the +# live inspection table for whichever single clip is on screen. + +def _difference_per_episode(episode_index, frame_index, values, fps): + """d(values)/dt within each episode, ordered by frame. values: (N,) or (N, d).""" + time = frame_index.astype(np.float64) / fps + rates = np.zeros(values.shape, dtype=np.float64) + for episode in np.unique(episode_index): + rows = np.where(episode_index == episode)[0] + rows = rows[np.argsort(frame_index[rows])] + if len(rows) < 2: + continue + rates[rows] = np.gradient(values[rows], time[rows], axis=0) + return rates + + +def add_temporal_features(features, episode_index, frame_index, fps): + """Add wrist speed, vertical velocity, and curl rate for the overlay table.""" + for tag in ("L", "R"): + wrist_velocity = _difference_per_episode(episode_index, frame_index, features[f"wrist_{tag}"], fps) + features[f"wrist_speed_{tag}"] = np.linalg.norm(wrist_velocity, axis=1) + features[f"wrist_vert_vel_{tag}"] = wrist_velocity[:, 1] + features[f"curl_rate_{tag}"] = _difference_per_episode(episode_index, frame_index, features[f"curl_{tag}"], fps) + return features + + +def add_angular_velocity(features, state, episode_index, frame_index, fps): + """Add wrist angular speed (rad/s) from consecutive orientations, for the overlay table.""" + for side, tag in HAND_TAGS: + rotations = rotation_from_rot6d(state[:, rot6d_slice(side)]) + angular_speed = np.zeros(len(rotations)) + for episode in np.unique(episode_index): + rows = np.where(episode_index == episode)[0] + rows = rows[np.argsort(frame_index[rows])] + if len(rows) < 2: + continue + relative = np.einsum("nij,nkj->nik", rotations[rows][1:], rotations[rows][:-1]) # R_t @ R_{t-1}^T + cosine = np.clip((np.trace(relative, axis1=1, axis2=2) - 1) / 2, -1, 1) + angular_speed[rows[1:]] = np.arccos(cosine) * fps + features[f"wrist_angvel_{tag}"] = angular_speed + return features diff --git a/datasets/egodex/query_ui.py b/datasets/egodex/query_ui.py new file mode 100644 index 0000000..1deddac --- /dev/null +++ b/datasets/egodex/query_ui.py @@ -0,0 +1,518 @@ +# /// script +# description = "Gradio UI: search EgoDex by hand-pose scenario plus SigLIP semantic text, with clip playback" +# requires-python = ">=3.10, <3.13" +# dependencies = ["daft>=0.7.15", "gradio>=5,<6", "torch", "transformers", "numpy"] +# /// +"""Pose-first scenario search over EgoDex, with semantic ranking and per-match playback. + +Primary control is the HAND-POSE query (top, always visible) — a Daft DataFrame +`.where(…)` over per-frame pose features. It defines *matching segments*: the +contiguous stretches where the predicate is true (the action's start→end). +The semantic text query (optional, below) only *ranks* the matched episodes. + +Flow: + • Search → master gallery of matched episodes (thumbnail at first match). + • Click an episode → detail: a LOOPING clip of just the matched segment, + its exact start/end, a segment stepper (for the N times it matches), and a + timeline showing where the matches sit in the episode. + + python query_ui.py # prints a public *.gradio.live URL +""" + +from __future__ import annotations + +import os +import json +import subprocess +import threading +import time +from collections import OrderedDict, defaultdict +from concurrent.futures import ThreadPoolExecutor +from itertools import count + +import daft +import gradio as gr +import numpy as np +import torch +from daft import col, lit +from transformers import AutoModel, AutoProcessor + +import lerobot +import pose_features as pf +from clip_features import DEVICE, MODEL_ID, _normalized_embedding + +HERE = os.path.dirname(os.path.abspath(__file__)) +OUT = os.environ.get("OUT", os.path.join(HERE, "out", "clip_features")) +POSE_OUT = os.environ.get("POSE_OUT", os.path.join(HERE, "out", "pose_features")) # precomputed geometry +DATASET = os.environ.get("DATASET", "./egodex_lerobot_full") # full pose(48+204)+video dataset +CLIPS_DIR, FRAMES_DIR = os.path.join(HERE, "ui_clips"), os.path.join(HERE, "ui_frames") +os.makedirs(CLIPS_DIR, exist_ok=True) +os.makedirs(FRAMES_DIR, exist_ok=True) + +KMAX = 24 # max episodes in the master gallery +SEG_GAP_MERGE = 5 # merge matching runs separated by < this many frames (~0.17s); keeps spans tight +SEG_MIN_FRAMES = 1 # keep even instant (1-frame) events, e.g. the moment of grasping +MAX_SEGS = 12 # cap segments shown per episode (keep the longest) +CLIP_PAD = 0.3 # seconds of lead-in / lead-out around a segment clip +MIN_CLIP_SECS = 1.2 # expand brief segments (e.g. a grasp) to at least this, so the loop is watchable +SEG_CLIP_MAX = 12.0 # cap clip length; a long continuous match shows its middle 12 s +SEMANTIC_WIN = 1.5 # seconds each side of the best frame for semantic-only (no pose) fallback +MAX_THUMBS = 16 # max exact-matched-frame thumbnails shown per segment +TWIST_ROLL_THR = 2.0 # rad/s of TRUE forearm roll (pronation/supination); absolute, not percentile + +# ── warm state: embeddings ────────────────────────────────────────────────── +print("loading embeddings…") +_e = daft.read_parquet(OUT).to_pydict() +EP = np.asarray(_e["episode_index"]); FR = np.asarray(_e["frame_index"]) +E = np.asarray(_e["clip_emb"], dtype=np.float32) +EP_ROWS = {int(x): np.where(EP == x)[0] for x in np.unique(EP)} +print(f" {E.shape[0]} embeddings over {len(EP_ROWS)} episodes") + +print("loading episode metadata…") +FPS = float(lerobot._read_info(lerobot._normalize_dataset_root(DATASET))["fps"]) +_k = "observation.image" +_m = lerobot.read_episodes(DATASET, include_video_metadata=True).select( + "episode_index", "tasks", f"videos/{_k}/chunk_index", f"videos/{_k}/file_index", + f"videos/{_k}/from_timestamp", f"videos/{_k}/to_timestamp").to_pydict() +TASK, WINDOW = {}, {} +for e, t, ci, fi, a, b in zip(_m["episode_index"], _m["tasks"], _m[f"videos/{_k}/chunk_index"], + _m[f"videos/{_k}/file_index"], _m[f"videos/{_k}/from_timestamp"], + _m[f"videos/{_k}/to_timestamp"]): + e = int(e) + TASK[e] = (t[0] if isinstance(t, (list, tuple)) and t else (t or "")) + WINDOW[e] = (os.path.join(DATASET, "videos", _k, f"chunk-{int(ci):03d}", f"file-{int(fi):03d}.mp4"), + float(a), float(b)) + +# ── warm state: raw pose, ONLY for the player's live skeleton overlay (viz) ── +print("loading raw pose for the player overlay…") +_p = lerobot.read(DATASET).where(col("episode_index").is_in(sorted(EP_ROWS))).select( + "episode_index", "frame_index", "observation.state", "observation.extrinsics").to_pydict() +pep = np.asarray(_p["episode_index"]); pfr = np.asarray(_p["frame_index"]) +S = np.asarray(_p["observation.state"], dtype=np.float32) +X = np.asarray(_p["observation.extrinsics"], dtype=np.float32) +# 48-D features for the live inspection table + landmark overlay ONLY (not the query path) +F = pf.compute_raw_features(S) +pf.add_temporal_features(F, pep, pfr, FPS) +pf.add_angular_velocity(F, S, pep, pfr, FPS) + +# ── query features: read the PRECOMPUTED geometric parquet (run_pose_features.py) ── +# Decoupled from the UI — scenarios are computed once offline, never recomputed here. +print("loading precomputed pose features…") +_q = daft.read_parquet(POSE_OUT).where(col("episode_index").is_in(sorted(EP_ROWS))).to_pydict() +qep = np.asarray(_q["episode_index"]); qfr = np.asarray(_q["frame_index"]) +# nearest embedded frame per query row (bridges 30 fps pose <-> 1 fps embeddings) +_by = defaultdict(list) +for i, (e, f) in enumerate(zip(EP, FR)): + _by[int(e)].append((int(f), i)) +_ep_emb = {e: (np.array([x[0] for x in sorted(v)]), np.array([x[1] for x in sorted(v)])) for e, v in _by.items()} +emb_row = np.empty(len(qep), dtype=np.int64) +for e in np.unique(qep): + idx = np.where(qep == e)[0]; ef, er = _ep_emb[int(e)] + pos = np.clip(np.searchsorted(ef, qfr[idx]), 0, len(ef) - 1) + prev = np.clip(pos - 1, 0, len(ef) - 1) + emb_row[idx] = er[np.where(np.abs(ef[pos] - qfr[idx]) <= np.abs(ef[prev] - qfr[idx]), pos, prev)] +POSE_DF = daft.from_pydict({**{k: np.asarray(v) for k, v in _q.items()}, "emb_row": emb_row}).collect() +_allclose = np.concatenate([np.asarray(_q["closure_L"]), np.asarray(_q["closure_R"])]) +CLOSE_LO, CLOSE_HI = float(np.percentile(_allclose, 2)), float(np.percentile(_allclose, 98)) +print(f" POSE_DF: {len(qep)} frames × precomputed pose features (read, not computed)") + +GRASP_RATE, LIFT_VEL = 0.20, 0.20 # 48-D thresholds: curl-closing rate (grasping), wrist up-speed (lifting) + +print("loading SigLIP-2 text tower…") +_model = AutoModel.from_pretrained(MODEL_ID).to(DEVICE).eval() +_proc = AutoProcessor.from_pretrained(MODEL_ID) + + +def encode(text): + inp = _proc(text=[text], return_tensors="pt", padding="max_length").to(DEVICE) + with torch.no_grad(): + return _normalized_embedding(_model.get_text_features(**inp))[0].cpu().numpy().astype(np.float32) + + +# ── media: thumbnails + looping segment clips, LRU-cached ──────────────────── +def _lru(maxsize): + cache, lock = OrderedDict(), threading.Lock() + + def put(key, path): + with lock: + cache[key] = path; cache.move_to_end(key) + while len(cache) > maxsize: + _, old = cache.popitem(last=False) + try: + os.remove(old) + except OSError: + pass + + def get(key): + with lock: + if key in cache: + cache.move_to_end(key); return cache[key] + return None + return get, put + + +_clip_get, _clip_put = _lru(300) +_frame_get, _frame_put = _lru(2000) + + +def get_frame(e, fr): + if (hit := _frame_get((e, fr))): + return hit + path = os.path.join(FRAMES_DIR, f"ep{e:04d}_f{fr:05d}.jpg") + if not os.path.exists(path): + shard, a, _ = WINDOW[e] + subprocess.run(["ffmpeg", "-y", "-ss", f"{a + fr / FPS:.4f}", "-i", shard, "-frames:v", "1", "-q:v", "3", + path, "-loglevel", "error"], check=True) + _frame_put((e, fr), path); return path + + +def get_segment_clip(e, start_frame, end_frame): + """A short, playable clip covering exactly [start_frame, end_frame] (+pad) — meant to be looped.""" + key = (e, int(start_frame), int(end_frame)) + if (hit := _clip_get(key)): + return hit + shard, a, b = WINDOW[e] + # center the clip on the segment, expanded to >= MIN_CLIP_SECS and capped at SEG_CLIP_MAX + want = min(SEG_CLIP_MAX, max(MIN_CLIP_SECS, (end_frame - start_frame) / FPS + 2 * CLIP_PAD)) + center = a + (start_frame + end_frame) / 2 / FPS + start_ts = max(a, center - want / 2) + dur = max(0.4, min(want, b - start_ts)) + path = os.path.join(CLIPS_DIR, f"ep{e:04d}_seg{int(start_frame):05d}_{int(end_frame):05d}.mp4") + if not os.path.exists(path): + # re-encode (not stream-copy) so the cut is FRAME-ACCURATE to [start,end] — we want just the match + subprocess.run(["ffmpeg", "-y", "-ss", f"{start_ts:.4f}", "-i", shard, "-t", f"{dur:.4f}", + "-c:v", "libx264", "-pix_fmt", "yuv420p", "-preset", "veryfast", "-crf", "23", + "-movflags", "+faststart", path, "-loglevel", "error"], check=True) + _clip_put(key, path); return path + + +# ── pose predicate (Daft) + segment extraction ────────────────────────────── +# Scenarios grouped by kind. States read one frame; actions read motion over frames. +STATE_SCENARIOS = ["any", "hand openness", "writing grip", "hammer grip"] +ACTION_SCENARIOS = ["any", "twisting", "lifting", "reaching", "grasping", "in-hand manipulation"] + + +def scenario_predicate(category, scenario, hand, open_lo, open_hi): + """Daft boolean for one selected scenario, or None (no pose filter).""" + if scenario in (None, "any"): + return None + # openness in [0,1] (1 = fully open palm) -> closure bounds (low closure = open) + span = (CLOSE_HI - CLOSE_LO) or 1.0 + clo_lo = CLOSE_HI - open_hi * span # higher openness -> lower closure + clo_hi = CLOSE_HI - open_lo * span + + def for_hand(h): + if scenario == "hand openness": # openness band, mapped onto the closure column + return (col(f"closure_{h}") >= clo_lo) & (col(f"closure_{h}") <= clo_hi) + if scenario == "writing grip": + return col(f"sc_writing_{h}") + if scenario == "hammer grip": + return col(f"sc_hammer_{h}") + if scenario == "twisting": + return col(f"sc_twisting_{h}") + if scenario == "reaching": + return col(f"sc_reaching_{h}") + if scenario == "in-hand manipulation": + return col(f"sc_inhand_{h}") + if scenario == "grasping": # action: 48-D curl closing + return col(f"curl_rate_{h}") <= -GRASP_RATE + if scenario == "lifting": # action: 48-D wrist moving up + return col(f"wrist_vert_vel_{h}") >= LIFT_VEL + return lit(True) + return for_hand("L") if hand == "left" else for_hand("R") if hand == "right" else (for_hand("L") | for_hand("R")) + + +def segments_of(frames): + """Contiguous matching runs (merging gaps < SEG_GAP_MERGE). Returns [(start, end), …].""" + frames = sorted(frames) + if not frames: + return [] + runs, s, prev = [], frames[0], frames[0] + for f in frames[1:]: + if f - prev > SEG_GAP_MERGE: + runs.append((s, prev)); s = f + prev = f + runs.append((s, prev)) + return [(a, b) for a, b in runs if b - a + 1 >= SEG_MIN_FRAMES] or runs + + +def search(query, k, hand, category, scenario, open_lo, open_hi): + """POSE_DF [join sim] → where(scenario predicate) → rank episodes → segment each.""" + pred = scenario_predicate(category, scenario, hand, open_lo, open_hi) + has_text = bool(query.strip()) + if pred is None and not has_text: + return [], gr.update(value="Pick a State/Action scenario (and/or type a Semantic query)."), [], "" + + df = POSE_DF + sims = None + if has_text: + sims = E @ encode(query.strip()) + df = df.join(daft.from_pydict({"emb_row": np.arange(len(sims)), "sim": sims.astype(np.float32)}), on="emb_row") + if pred is not None: + df = df.where(pred) + score = (col("sim").max() if has_text else col("frame_index").count()).alias("score") + ranked = df.groupby("episode_index").agg(col("frame_index").count().alias("n"), score).sort("score", desc=True).limit(int(k)).to_pydict() + top = [int(e) for e in ranked["episode_index"]] + scores = list(ranked["score"]) + + # matching frames for the ranked episodes → per-episode segments + mf = df.where(col("episode_index").is_in(top)).select("episode_index", "frame_index").to_pydict() + frames_by_ep = defaultdict(list) + for e, f in zip(mf["episode_index"], mf["frame_index"]): + frames_by_ep[int(e)].append(int(f)) + + results = [] + for e, sc in zip(top, scores): + if pred is not None: + segs = sorted(sorted(segments_of(frames_by_ep[e]), key=lambda p: p[1] - p[0], reverse=True)[:MAX_SEGS]) + framelist = sorted(frames_by_ep[e]) # exact frames where the predicate is true + else: # semantic-only: a window around the episode's best-matching frame + j = EP_ROWS[e][int(np.argmax(sims[EP_ROWS[e]]))] + c = int(FR[j]); w = int(SEMANTIC_WIN * FPS) + segs = [(max(0, c - w), c + w)] + framelist = list(range(segs[0][0], segs[0][1] + 1, max(1, int(FPS // 3)))) + if segs: + results.append({"ep": e, "score": float(sc), "has_text": has_text, "segs": segs, "frames": framelist, "task": TASK.get(e, "")}) + + items = [] + for r in results: + s0 = r["segs"][0][0] + sct = f"sim {r['score']:.3f}" if has_text else f"{int(r['score'])} frames" + items.append((get_frame(r["ep"], s0), f"ep {r['ep']} · {sct} · {len(r['segs'])} match(es)")) + status = f"✅ {len(results)} episodes" + (f" · ranked by {query.strip()!r}" if has_text else "") + " · click one (left) to inspect its matches" + return items, gr.update(value=status), results, "" + + +# ── client-side player: per-frame pose pushed to the browser, synced to