Add 0.5B Alpamayo validation
This commit is contained in:
seeker
@@ -0,0 +1,204 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from huggingface_hub import get_token, hf_hub_download
|
||||
from physical_ai_av import PhysicalAIAVDatasetInterface
|
||||
from PIL import Image
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
from alpamayo_lite.config import load_config
|
||||
from alpamayo_lite.prompts import clean_reasoning_text
|
||||
|
||||
|
||||
CLIP_DURATION_US = 20_000_000
|
||||
|
||||
|
||||
def parse_events(cell: object) -> list[dict[str, Any]]:
|
||||
if cell is None or (np.isscalar(cell) and pd.isna(cell)):
|
||||
return []
|
||||
parsed = json.loads(cell) if isinstance(cell, str) else cell
|
||||
if not isinstance(parsed, (list, tuple)):
|
||||
return []
|
||||
return [event for event in parsed if isinstance(event, dict)]
|
||||
|
||||
|
||||
def collect_candidates(
|
||||
reasoning_path: Path,
|
||||
history_seconds: float,
|
||||
future_seconds: float,
|
||||
) -> list[dict[str, Any]]:
|
||||
frame = pd.read_parquet(reasoning_path)
|
||||
candidates: list[dict[str, Any]] = []
|
||||
for clip_id, events_cell in frame["events"].items():
|
||||
for event in parse_events(events_cell):
|
||||
if "event_start_timestamp" not in event:
|
||||
continue
|
||||
t0_us = int(event["event_start_timestamp"])
|
||||
if t0_us < int(history_seconds * 1_000_000):
|
||||
continue
|
||||
if t0_us + int(future_seconds * 1_000_000) > CLIP_DURATION_US:
|
||||
continue
|
||||
# Current PAI releases call the driving rationale/command `coc`.
|
||||
# Keep `cot` as a compatibility fallback for older metadata shards.
|
||||
reasoning = clean_reasoning_text(str(event.get("coc") or event.get("cot") or ""))
|
||||
candidates.append({"clip_id": str(clip_id), "t0_us": t0_us, "reasoning": reasoning})
|
||||
if not candidates:
|
||||
raise RuntimeError("No valid reasoning events remain after trajectory safety filtering")
|
||||
return candidates
|
||||
|
||||
|
||||
def load_local_trajectory(
|
||||
interface: PhysicalAIAVDatasetInterface,
|
||||
clip_id: str,
|
||||
t0_us: int,
|
||||
history_steps: int,
|
||||
future_steps: int,
|
||||
time_step: float,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
egomotion = interface.get_clip_feature(
|
||||
clip_id, interface.features.LABELS.EGOMOTION, maybe_stream=True
|
||||
)
|
||||
history_offsets = np.arange(-(history_steps - 1), 1, dtype=np.int64)
|
||||
history_timestamps = t0_us + history_offsets * int(time_step * 1_000_000)
|
||||
future_offsets = np.arange(1, future_steps + 1, dtype=np.int64)
|
||||
future_timestamps = t0_us + future_offsets * int(time_step * 1_000_000)
|
||||
|
||||
history_state = egomotion(history_timestamps)
|
||||
future_state = egomotion(future_timestamps)
|
||||
history_xyz = history_state.pose.translation
|
||||
future_xyz = future_state.pose.translation
|
||||
t0_rotation = Rotation.from_quat(history_state.pose.rotation.as_quat()[-1])
|
||||
inverse = t0_rotation.inv()
|
||||
origin = history_xyz[-1].copy()
|
||||
history_local = inverse.apply(history_xyz - origin).astype(np.float32)
|
||||
future_local = inverse.apply(future_xyz - origin).astype(np.float32)
|
||||
return history_local, future_local
|
||||
|
||||
|
||||
def load_camera_frame(
|
||||
interface: PhysicalAIAVDatasetInterface,
|
||||
clip_id: str,
|
||||
t0_us: int,
|
||||
camera_name: str,
|
||||
) -> np.ndarray:
|
||||
camera_feature = getattr(interface.features.CAMERA, camera_name.upper())
|
||||
reader = interface.get_clip_feature(clip_id, camera_feature, maybe_stream=True)
|
||||
frames, _ = reader.decode_images_from_timestamps(np.asarray([t0_us], dtype=np.int64))
|
||||
return frames[0]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--config", default="configs/local_overfit.yaml")
|
||||
args = parser.parse_args()
|
||||
cfg = load_config(args.config)
|
||||
token = os.environ.get("HF_TOKEN") or get_token()
|
||||
if not token:
|
||||
raise SystemExit(
|
||||
"PAI is gated. Set HF_TOKEN for an account that accepted the dataset license."
|
||||
)
|
||||
|
||||
raw_dir = cfg.resolve(cfg.paths.pai_raw_dir)
|
||||
raw_dir.mkdir(parents=True, exist_ok=True)
|
||||
cache_dir = cfg.root / ".cache" / "huggingface"
|
||||
reasoning_file = Path(
|
||||
hf_hub_download(
|
||||
repo_id=cfg.data.repo_id,
|
||||
repo_type="dataset",
|
||||
filename="reasoning/ood_reasoning.parquet",
|
||||
token=token,
|
||||
cache_dir=cache_dir,
|
||||
local_dir=raw_dir,
|
||||
)
|
||||
)
|
||||
candidates = collect_candidates(
|
||||
reasoning_file,
|
||||
history_seconds=(cfg.data.history_steps - 1) * cfg.data.time_step,
|
||||
future_seconds=cfg.data.source_future_steps * cfg.data.time_step,
|
||||
)
|
||||
rng = np.random.default_rng(cfg.seed)
|
||||
selected_indices = rng.choice(
|
||||
len(candidates), size=min(cfg.data.num_clips, len(candidates)), replace=False
|
||||
)
|
||||
selected = [candidates[int(index)] for index in selected_indices]
|
||||
|
||||
interface = PhysicalAIAVDatasetInterface(
|
||||
token=token,
|
||||
cache_dir=cache_dir,
|
||||
local_dir=raw_dir,
|
||||
confirm_download_threshold_gb=1000.0,
|
||||
)
|
||||
processed_dir = cfg.resolve(cfg.paths.manifest).parent
|
||||
image_dir = processed_dir / "images"
|
||||
trajectory_dir = processed_dir / "trajectories"
|
||||
image_dir.mkdir(parents=True, exist_ok=True)
|
||||
trajectory_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
rows: list[dict[str, Any]] = []
|
||||
for sample_index, event in enumerate(selected):
|
||||
clip_id = event["clip_id"]
|
||||
t0_us = event["t0_us"]
|
||||
history, future_full = load_local_trajectory(
|
||||
interface=interface,
|
||||
clip_id=clip_id,
|
||||
t0_us=t0_us,
|
||||
history_steps=cfg.data.history_steps,
|
||||
future_steps=cfg.data.source_future_steps,
|
||||
time_step=cfg.data.time_step,
|
||||
)
|
||||
waypoint_indices = np.linspace(
|
||||
0,
|
||||
cfg.data.source_future_steps - 1,
|
||||
cfg.data.expert_waypoints,
|
||||
dtype=np.int64,
|
||||
)
|
||||
future = future_full[waypoint_indices]
|
||||
frame = load_camera_frame(
|
||||
interface=interface,
|
||||
clip_id=clip_id,
|
||||
t0_us=t0_us,
|
||||
camera_name=cfg.data.camera,
|
||||
)
|
||||
|
||||
stem = f"sample_{sample_index:02d}_{clip_id}_{t0_us}"
|
||||
image_path = image_dir / f"{stem}.jpg"
|
||||
trajectory_path = trajectory_dir / f"{stem}.npz"
|
||||
Image.fromarray(frame).save(image_path, quality=95)
|
||||
np.savez_compressed(
|
||||
trajectory_path,
|
||||
history_xyz=history,
|
||||
future_xyz=future,
|
||||
source_future_xyz=future_full,
|
||||
waypoint_indices=waypoint_indices,
|
||||
)
|
||||
rows.append(
|
||||
{
|
||||
"clip_id": clip_id,
|
||||
"t0_us": t0_us,
|
||||
"reasoning": event["reasoning"],
|
||||
"image_path": str(image_path.relative_to(cfg.root)),
|
||||
"trajectory_path": str(trajectory_path.relative_to(cfg.root)),
|
||||
"camera": cfg.data.camera,
|
||||
"history_steps": cfg.data.history_steps,
|
||||
"expert_waypoints": cfg.data.expert_waypoints,
|
||||
}
|
||||
)
|
||||
|
||||
manifest_path = cfg.resolve(cfg.paths.manifest)
|
||||
manifest_path.write_text(
|
||||
"".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
|
||||
encoding="utf-8",
|
||||
)
|
||||
print(json.dumps({"manifest": str(manifest_path), "samples": rows}, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user