File size: 6,261 Bytes
319eb16 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | #!/usr/bin/env python3
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()
|