| from pathlib import Path |
| import argparse |
| import json |
| import pickle |
| import sys |
|
|
| import numpy as np |
|
|
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| if str(PROJECT_ROOT) not in sys.path: |
| sys.path.insert(0, str(PROJECT_ROOT)) |
|
|
| from rr_label_study.oven_study import ( |
| BimanualTakeTrayOutOfOven, |
| ReplayCache, |
| _launch_replay_env, |
| _load_demo, |
| _pregrasp_progress_and_distance, |
| _pregrasp_score_and_success, |
| ) |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--episode-dir", required=True) |
| parser.add_argument("--templates-pkl", required=True) |
| parser.add_argument("--frame-indices", nargs="+", type=int, required=True) |
| parser.add_argument("--checkpoint-stride", type=int, default=16) |
| parser.add_argument("--output-dir", required=True) |
| args = parser.parse_args() |
|
|
| output_dir = Path(args.output_dir) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| frame_indices = sorted(set(args.frame_indices)) |
| pending_frame_indices = [ |
| frame_index |
| for frame_index in frame_indices |
| if not output_dir.joinpath(f"frame_{frame_index:04d}.json").exists() |
| ] |
| if not pending_frame_indices: |
| return 0 |
|
|
| episode_dir = Path(args.episode_dir) |
| with Path(args.templates_pkl).open("rb") as handle: |
| templates = pickle.load(handle) |
| demo = _load_demo(episode_dir) |
|
|
| env = _launch_replay_env() |
| try: |
| task = env.get_task(BimanualTakeTrayOutOfOven) |
| cache = ReplayCache(task, demo, checkpoint_stride=args.checkpoint_stride) |
| cache.reset() |
| total = len(pending_frame_indices) |
| for completed, frame_index in enumerate(pending_frame_indices, start=1): |
| cache.step_to(frame_index) |
| state = cache.current_state() |
| pregrasp_progress, pregrasp_distance = _pregrasp_progress_and_distance( |
| np.asarray(state.left_gripper_pose, dtype=np.float64), |
| np.asarray(state.tray_pose, dtype=np.float64), |
| templates, |
| ) |
| p_pre, y_pre, _ = _pregrasp_score_and_success(task, templates) |
| row = { |
| "frame_index": int(frame_index), |
| "pregrasp_progress": float(pregrasp_progress), |
| "pregrasp_distance": float(pregrasp_distance), |
| "p_pre": float(p_pre), |
| "y_pre_raw": float(bool(y_pre)), |
| "y_pre": float(bool(y_pre)), |
| } |
| row_path = output_dir.joinpath(f"frame_{frame_index:04d}.json") |
| tmp_path = row_path.with_suffix(".json.tmp") |
| with tmp_path.open("w", encoding="utf-8") as handle: |
| json.dump(row, handle) |
| tmp_path.replace(row_path) |
| if completed == total or completed % 8 == 0: |
| print( |
| json.dumps( |
| { |
| "done": completed, |
| "total": total, |
| "frame_index": int(frame_index), |
| } |
| ), |
| flush=True, |
| ) |
| finally: |
| env.shutdown() |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
