extract frame

indexing/tests/e2e.py

@@ -3,14 +3,20 @@ End-to-end test: pulls images and videos from HF datasets, indexes them
 with ragstudio, syncs the index to the team-0/ragstudio bucket, then runs
 a sample search over each modality.
 
+Both the downloaded datasets and the FAISS index live under tests/.cache/,
+so re-runs reuse them. Pass --force to rebuild the index from scratch.
+
 Usage:
-    python tests/e2e.py
+    python tests/e2e.py [--force]
 """
 
+import argparse
 import os
+import re
 import sys
 from pathlib import Path
 
+import cv2
 from huggingface_hub import snapshot_download
 
 ROOT = Path(__file__).resolve().parent.parent
@@ -21,10 +27,42 @@ VIDEO_DATASET = "team-0/test-videos"
 
 IMAGE_QUERY = "a person presenting on stage"
 
-DATA_DIR = ROOT / "tests" / ".cache" / "data"
+CACHE_DIR = ROOT / "tests" / ".cache"
+DATA_DIR = CACHE_DIR / "data"
+FRAMES_DIR = CACHE_DIR / "frames"
+
+# Video meta entries are formatted by indexers/video.py as "<path> @ <t>s".
+VIDEO_META_RE = re.compile(r"^(.*) @ (\d+\.\d+)s$")
+
+
+def extract_video_frame(meta: str, out_dir: Path) -> Path:
+    m = VIDEO_META_RE.match(meta)
+    if not m:
+        raise ValueError(f"unexpected video meta format: {meta!r}")
+    vpath, t = m.group(1), float(m.group(2))
+
+    cap = cv2.VideoCapture(vpath)
+    fps = cap.get(cv2.CAP_PROP_FPS) or 0
+    if fps > 0:
+        cap.set(cv2.CAP_PROP_POS_FRAMES, int(round(t * fps)))
+    else:
+        cap.set(cv2.CAP_PROP_POS_MSEC, t * 1000)
+    ok, frame = cap.read()
+    cap.release()
+    if not ok:
+        raise RuntimeError(f"failed to read frame at {t}s from {vpath}")
+
+    out_dir.mkdir(parents=True, exist_ok=True)
+    out_path = out_dir / f"{Path(vpath).stem}_{t:.2f}s.png"
+    cv2.imwrite(str(out_path), frame)
+    return out_path
 
 
 def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.strip())
+    parser.add_argument("--force", action="store_true", help="rebuild the index even if cached")
+    args = parser.parse_args()
+
     DATA_DIR.mkdir(parents=True, exist_ok=True)
 
     print(f"[1/4] Downloading datasets to {DATA_DIR}")
@@ -51,10 +89,10 @@ def main() -> None:
     # at <repo>/index_data/ instead of vanishing with a temp dir.
     os.chdir(ROOT)
 
-    print(f"[2/4] Building index from {DATA_DIR}")
+    print(f"[2/4] Building index from {DATA_DIR}{' (force)' if args.force else ''}")
     from index import build_index
 
-    build_index(DATA_DIR)
+    build_index(DATA_DIR, force=args.force)
 
     print("[3/4] Syncing index to team-0/ragstudio bucket")
     from sync import sync
@@ -63,6 +101,7 @@ def main() -> None:
     print("[4/4] Searching")
     from searchers import SEARCHERS
 
+
     image_hits = SEARCHERS["image"](IMAGE_QUERY, top_k=3)
     print(f"\n=== image: {IMAGE_QUERY!r} ===")
     for s, p in image_hits:
@@ -72,10 +111,12 @@ def main() -> None:
     for q in video_queries:
         video_hits = SEARCHERS["video"](q, top_k=3)
         print(f"\n=== video: {q!r} ===")
-        for s, p in video_hits:
-            print(f"  {s:.3f}  {p}")
+        for s, meta in video_hits:
+            frame_path = extract_video_frame(meta, FRAMES_DIR)
+            print(f"  {s:.3f}  {meta}  ->  {frame_path}")
         assert video_hits, f"no video results for {q!r}"
 
+
     print("\nOK")