File size: 4,006 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 | #!/usr/bin/env python3
"""Run Robometer progress scoring on the BenchVideo clips (no marks).
For each episode in a category's tasks.json:
- load mp4 -> sample to <= max_frames at fps
- frame_steps inference -> per-frame progress curve
- save results_robometer.json + overview_robometer.png
Output: BenchVideo/eval_results/<Category>/<id>/
"""
import argparse
import json
import os
from pathlib import Path
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from benchmark_progress_mark_local import (
RobometerLocalRunner,
load_video_frames_with_indices,
)
BENCH_ROOT = Path("/home/vcj9002/jianshu/workspace/code_keliang/Current_Baseline/Benchmark_Videos/BenchVideo")
EVAL_ROOT = BENCH_ROOT / "eval_results"
MODEL_PATH = "/home/vcj9002/jianshu/workspace/code_keliang/Current_Baseline/Robometer/models/Robometer-4B"
def write_overview(out_path, frames, progress, title):
fig, ax = plt.subplots(figsize=(11, 4.6))
ax.plot(frames[: len(progress)], progress, marker="o", linewidth=1.5, markersize=2.8, color="#b9482f")
ax.set_ylim(-0.05, 1.05)
ax.set_xlim(0, max(5, int(max(frames) * 1.02)) if frames else 5)
ax.set_xlabel("Original Frame"); ax.set_ylabel("Progress (normalized)")
ax.set_title(title); ax.grid(True, linestyle="--", alpha=0.35)
fig.tight_layout(); fig.savefig(out_path, dpi=180); plt.close(fig)
def run_category(runner, category, fps, max_frames, fallback):
cat_dir = BENCH_ROOT / category
data = json.loads((cat_dir / "tasks.json").read_text(encoding="utf-8"))
out_root = EVAL_ROOT / category
out_root.mkdir(parents=True, exist_ok=True)
for ep in data["episodes"]:
name = ep["episode"]; stem = name[:-4]
task = (ep.get("tasks") or [fallback])[0] or fallback
src = cat_dir / name
ep_out = out_root / stem
if (ep_out / "results_robometer.json").exists():
print(f"[skip] {category}/{name} (results_robometer.json exists)")
continue
ep_out.mkdir(parents=True, exist_ok=True)
print(f"\n[RUN] {category}/{name} task={task!r}")
frames, sampled_idx, total, native_fps = load_video_frames_with_indices(
src, fps=fps, max_frames=max_frames, required_frames=[])
progress, success = runner.compute_rewards_per_frame(
frames, task, inference_mode="frame_steps",
prefix_sample_frames=8, prefix_batch_size=1)
progress = [float(x) for x in progress]
results = {
"category": category, "episode": name,
"original_file": ep.get("original_file"), "task": task,
"num_total_frames": total, "native_fps": native_fps, "sample_fps": fps,
"num_sampled_frames": len(sampled_idx),
"sampled_original_frame_indices": sampled_idx,
"progress_curve": progress,
}
(ep_out / "results_robometer.json").write_text(
json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
write_overview(ep_out / "overview_robometer.png", sampled_idx, progress,
f"Robometer {category}/{stem}: {task[:50]}")
rng = f"{min(progress):.2f}~{max(progress):.2f}" if progress else "empty"
print(f"[OK] {ep_out} (progress {rng})")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--category", default="all")
ap.add_argument("--fps", type=float, default=2.0)
ap.add_argument("--max-frames", type=int, default=128)
ap.add_argument("--instruction-fallback", default="Complete the task shown in the video.")
args = ap.parse_args()
runner = RobometerLocalRunner(model_path=MODEL_PATH)
cats = ["Count", "State", "Sequence"] if args.category == "all" else [args.category]
for c in cats:
run_category(runner, c, args.fps, args.max_frames, args.instruction_fallback)
if __name__ == "__main__":
main()
|