| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from benchmark_progress_mark_local import ( |
| DEFAULT_INFERENCE_MODE, |
| DEFAULT_MAX_FRAMES, |
| DEFAULT_MIN_FRAMES, |
| DEFAULT_MODEL_PATH, |
| DEFAULT_PREFIX_BATCH_SIZE, |
| DEFAULT_PREFIX_SAMPLE_FRAMES, |
| RobometerLocalRunner, |
| build_frame_retry_schedule, |
| is_cuda_oom_error, |
| load_all_video_frames, |
| sample_video_frames_with_indices, |
| write_overview_plot, |
| ) |
|
|
|
|
| def save_results(output_dir: Path, payload: dict) -> None: |
| with (output_dir / "results.json").open("w", encoding="utf-8") as f: |
| json.dump(payload, f, indent=2, ensure_ascii=False) |
| f.write("\n") |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser( |
| description="Run Robometer on a single video with the same benchmark-style inference logic, without progress marks." |
| ) |
| parser.add_argument("--video", required=True, help="Path to the input video") |
| parser.add_argument("--task", required=True, help="Task instruction") |
| parser.add_argument("--model-path", default=str(DEFAULT_MODEL_PATH)) |
| parser.add_argument("--fps", type=float, default=3.0) |
| parser.add_argument("--max-frames", type=int, default=DEFAULT_MAX_FRAMES) |
| parser.add_argument("--min-frames", type=int, default=DEFAULT_MIN_FRAMES) |
| parser.add_argument( |
| "--inference-mode", |
| choices=["frame_steps", "whole"], |
| default=DEFAULT_INFERENCE_MODE, |
| help="frame_steps matches benchmark behavior; whole does one forward pass on the full sampled trajectory", |
| ) |
| parser.add_argument("--prefix-sample-frames", type=int, default=DEFAULT_PREFIX_SAMPLE_FRAMES) |
| parser.add_argument("--prefix-batch-size", type=int, default=DEFAULT_PREFIX_BATCH_SIZE) |
| parser.add_argument( |
| "--adaptive-max-frames", |
| dest="adaptive_max_frames", |
| action="store_true", |
| help="On CUDA OOM in whole mode, retry with smaller frame budgets", |
| ) |
| parser.add_argument( |
| "--no-adaptive-max-frames", |
| dest="adaptive_max_frames", |
| action="store_false", |
| help="Disable frame-budget retry on CUDA OOM", |
| ) |
| parser.add_argument( |
| "--output-root", |
| type=Path, |
| default=Path(__file__).parent / "outputs_benchmark_style", |
| help="Root directory for outputs", |
| ) |
| parser.set_defaults(adaptive_max_frames=True) |
| args = parser.parse_args() |
|
|
| video_path = Path(args.video).expanduser().resolve() |
| if not video_path.exists(): |
| raise FileNotFoundError(f"Video not found: {video_path}") |
|
|
| output_dir = args.output_root / video_path.stem |
| output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| runner = RobometerLocalRunner(model_path=str(args.model_path)) |
|
|
| print(f"[RUN] {video_path.name}") |
| print(f"[TASK] {args.task}") |
| print(f"[VIDEO] {video_path}") |
| print(f"[OUT] {output_dir}") |
|
|
| all_frames, native_fps = load_all_video_frames(video_path) |
| total_frames = len(all_frames) |
| retry_schedule = ( |
| build_frame_retry_schedule(args.max_frames, args.min_frames, bool(args.adaptive_max_frames)) |
| if args.inference_mode == "whole" |
| else [int(args.max_frames)] |
| ) |
|
|
| progress_pred = None |
| success_probs = None |
| sampled_indices = None |
| used_max_frames = retry_schedule[0] |
|
|
| for attempt_idx, frame_budget in enumerate(retry_schedule, start=1): |
| frames, sampled_indices = sample_video_frames_with_indices( |
| all_frames, |
| native_fps=native_fps, |
| fps=float(args.fps), |
| max_frames=int(frame_budget), |
| required_frames=[], |
| ) |
| print( |
| f"Loaded {total_frames} total frames; sampled {len(frames)} frames at fps={float(args.fps):g} " |
| f"(max_frames={int(frame_budget)}, try {attempt_idx}/{len(retry_schedule)})" |
| ) |
| try: |
| progress_pred, success_probs = runner.compute_rewards_per_frame( |
| video_frames=frames, |
| task=args.task, |
| inference_mode=args.inference_mode, |
| prefix_sample_frames=int(args.prefix_sample_frames), |
| prefix_batch_size=int(args.prefix_batch_size), |
| ) |
| used_max_frames = int(frame_budget) |
| break |
| except RuntimeError as exc: |
| if args.inference_mode != "whole" or not is_cuda_oom_error(exc) or attempt_idx == len(retry_schedule): |
| raise |
| next_budget = retry_schedule[attempt_idx] |
| print( |
| f"[OOM] whole inference hit CUDA OOM at max_frames={int(frame_budget)}; " |
| f"retrying with max_frames={int(next_budget)}" |
| ) |
| runner.reload_model() |
|
|
| if progress_pred is None or success_probs is None or sampled_indices is None: |
| raise RuntimeError("Failed to compute Robometer outputs.") |
| if progress_pred.size == 0: |
| raise RuntimeError("Robometer returned empty progress predictions.") |
| if progress_pred.size != len(sampled_indices): |
| raise RuntimeError( |
| f"Progress length mismatch: got {progress_pred.size} predictions for {len(sampled_indices)} sampled frames" |
| ) |
|
|
| np.save(str(output_dir / "progress.npy"), progress_pred) |
| np.save(str(output_dir / "success_probs.npy"), success_probs) |
|
|
| payload = { |
| "video": str(video_path), |
| "instruction": args.task, |
| "num_total_frames": total_frames, |
| "native_fps": native_fps, |
| "sample_fps": float(args.fps), |
| "num_sampled_frames": len(sampled_indices), |
| "sampled_original_frame_indices": sampled_indices, |
| "max_frames_used": used_max_frames, |
| "inference_mode": args.inference_mode, |
| "prefix_sample_frames": int(args.prefix_sample_frames), |
| "prefix_batch_size": int(args.prefix_batch_size), |
| "progress_min": float(np.min(progress_pred)), |
| "progress_max": float(np.max(progress_pred)), |
| "progress_mean": float(np.mean(progress_pred)), |
| } |
| save_results(output_dir, payload) |
| write_overview_plot(output_dir, sampled_indices, progress_pred, []) |
| print(f"[OK] Saved results to {output_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
