| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| from PIL import Image |
|
|
|
|
| def main() -> None: |
| t0 = time.time() |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--checkpoint-dir", required=True, help="Local checkpoint directory, e.g. .../4000") |
| parser.add_argument("--session-root", default="/workspace/data/teddybear_raw/session_20260327_165944_bear") |
| parser.add_argument("--sync-row-index", type=int, default=0) |
| parser.add_argument("--config-name", default="pi05_kinova_teddybear") |
| parser.add_argument("--output-json", default="/workspace/kinova_scene_sim/outputs/pi_actions.json") |
| parser.add_argument("--prompt", default="pick up the teddy bear and place it in the red box") |
| args = parser.parse_args() |
|
|
| session_root = Path(args.session_root) |
| print(f"[export] loading session from {session_root}", flush=True) |
| sync = pd.read_csv(session_root / "sync_index.csv") |
| row = sync.iloc[args.sync_row_index] |
|
|
| azure_rgb = np.array(Image.open(session_root / row["azure_rgb_file"]).convert("RGB")) |
| wrist_rgb = np.array(Image.open(session_root / row["rgb_file"]).convert("RGB")) |
| state = np.array( |
| [ |
| row["tool_x_m"], |
| row["tool_y_m"], |
| row["tool_z_m"], |
| row["tool_theta_x_deg"], |
| row["tool_theta_y_deg"], |
| row["tool_theta_z_deg"], |
| row["gripper_pos"], |
| ], |
| dtype=np.float32, |
| ) |
| print(f"[export] loaded observation in {time.time() - t0:.2f}s", flush=True) |
|
|
| print("[export] importing openpi modules", flush=True) |
| from openpi.policies import policy_config as _policy_config |
| from openpi.training import config as _config |
|
|
| print(f"[export] imported openpi in {time.time() - t0:.2f}s", flush=True) |
| print(f"[export] building train config {args.config_name}", flush=True) |
| train_config = _config.get_config(args.config_name) |
| print(f"[export] creating policy from {args.checkpoint_dir}", flush=True) |
| policy = _policy_config.create_trained_policy(train_config, Path(args.checkpoint_dir)) |
| print(f"[export] policy ready in {time.time() - t0:.2f}s", flush=True) |
| print("[export] running inference", flush=True) |
| results = policy.infer( |
| { |
| "observation/state": state, |
| "observation/image": azure_rgb, |
| "observation/wrist_image": wrist_rgb, |
| "prompt": args.prompt, |
| } |
| ) |
| print(f"[export] inference done in {time.time() - t0:.2f}s", flush=True) |
|
|
| actions = np.asarray(results["actions"], dtype=np.float32) |
| output = { |
| "checkpoint_dir": str(Path(args.checkpoint_dir)), |
| "session_root": str(session_root), |
| "sync_row_index": int(args.sync_row_index), |
| "azure_rgb_file": str(row["azure_rgb_file"]), |
| "wrist_rgb_file": str(row["rgb_file"]), |
| "state": state.tolist(), |
| "actions": actions.tolist(), |
| } |
|
|
| output_path = Path(args.output_json) |
| output_path.parent.mkdir(parents=True, exist_ok=True) |
| output_path.write_text(json.dumps(output, indent=2)) |
| print(f"saved action chunk to {output_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
