Buckets:
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import h5py | |
| import numpy as np | |
| from remotezip import RemoteZip | |
| from . import paths | |
| from .geometry import Z_UP_FROM_ARKIT, to_z_up | |
| from .hands import Hand | |
| from .skeleton import ARKIT, SIDES, joint_name | |
| BASE_URL = "https://ml-site.cdn-apple.com/datasets/egodex" | |
| FPS = 30.0 | |
| def fetch(task, index, split="test"): | |
| folder = paths.DATA / "egodex" / split / task | |
| targets = [folder / f"{index}.hdf5", folder / f"{index}.mp4"] | |
| if not all(target.exists() for target in targets): | |
| with RemoteZip(f"{BASE_URL}/{split}.zip") as archive: | |
| for suffix in ("hdf5", "mp4"): | |
| archive.extract( | |
| f"{split}/{task}/{index}.{suffix}", paths.DATA / "egodex" | |
| ) | |
| return targets[0] | |
| def _description(attrs): | |
| which = int(attrs.get("which_llm_description", 1)) | |
| key = "llm_description2" if which == 2 else "llm_description" | |
| return str(attrs.get(key, attrs.get("description", ""))) | |
| class Episode: | |
| task: str | |
| index: int | |
| text: str | |
| fps: float | |
| intrinsics: np.ndarray | |
| camera: np.ndarray | |
| hands: dict | |
| video_path: Path | |
| def __len__(self): | |
| return len(self.camera) | |
| def load(cls, task, index, split="test"): | |
| hdf5_path = fetch(task, index, split) | |
| with h5py.File(hdf5_path) as f: | |
| camera = to_z_up(f["transforms/camera"][:]) | |
| hands = {} | |
| for side in SIDES: | |
| names = [joint_name(side, joint) for joint in ARKIT.joints] | |
| poses = to_z_up( | |
| np.stack([f[f"transforms/{name}"][:] for name in names], axis=1) | |
| ) | |
| confidence = np.stack( | |
| [ | |
| f[f"confidences/{name}"][:] | |
| if "confidences" in f | |
| else np.ones(len(poses)) | |
| for name in names | |
| ], | |
| axis=1, | |
| ) | |
| assert poses.shape[1] == ARKIT.size | |
| hands[side] = Hand(side, poses, confidence) | |
| return cls( | |
| task=task, | |
| index=index, | |
| text=_description(f.attrs), | |
| fps=FPS, | |
| intrinsics=f["camera/intrinsic"][:], | |
| camera=camera, | |
| hands=hands, | |
| video_path=hdf5_path.with_suffix(".mp4"), | |
| ) | |
| def save(self, path): | |
| arrays = { | |
| "meta/task": self.task, | |
| "meta/index": self.index, | |
| "meta/text": self.text, | |
| "meta/fps": self.fps, | |
| "meta/source": "egodex", | |
| "meta/frame": "world_zup_m", | |
| "T_world_arkit": np.eye(4) * 1.0, | |
| "camera/K": self.intrinsics, | |
| "camera/T_world_cam": self.camera, | |
| } | |
| arrays["T_world_arkit"][:3, :3] = Z_UP_FROM_ARKIT | |
| for side, hand in self.hands.items(): | |
| arrays[f"hands/{side}/joint_SE3"] = hand.poses | |
| arrays[f"hands/{side}/kp21"] = hand.keypoints | |
| arrays[f"hands/{side}/conf"] = hand.confidence | |
| np.savez_compressed(path, **arrays) | |
Xet Storage Details
- Size:
- 3.21 kB
- Xet hash:
- 9a35dfe1c1ba4c047f814b9cbc17e94a1cc605e2f8bbf63142ab67cf1efef195
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.