Spaces:
Sleeping
Sleeping
File size: 4,848 Bytes
f6a6455 | 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 | """Stage 5 + 6: candidate frames and the weighted screenshot indicator.
Flow:
* ``extract_candidates`` pulls a dense, downscaled grid of frames from the video in a
single ffmpeg pass, *before* the video is deleted.
* ``delete_video`` removes the source video (auto-delete once the transcript + frames
exist).
* ``select_screenshots`` runs *after* the LLM. For each step it computes a capture time
that blends the LLM's suggested timestamp with the Whisper segment timing of the text
the step quotes, then snaps that target to the nearest pre-extracted candidate frame.
"""
from __future__ import annotations
import difflib
import glob
import os
import subprocess
# A 1 fps grid gives <=0.5 s snapping error. Capped so a long video can't explode disk.
GRID_FPS = 1.0
MAX_CANDIDATES = 1500
def extract_candidates(video_path: str, out_dir: str, duration: float) -> list[tuple[float, str]]:
"""Extract a downscaled candidate-frame grid; return sorted ``[(time_s, path), ...]``."""
os.makedirs(out_dir, exist_ok=True)
eff_fps = GRID_FPS
if duration and duration * GRID_FPS > MAX_CANDIDATES:
eff_fps = MAX_CANDIDATES / duration # thin the grid for very long videos
pattern = os.path.join(out_dir, "cand_%06d.jpg")
cmd = [
"ffmpeg", "-y", "-i", video_path,
"-vf", f"fps={eff_fps:.6f},scale=-2:720",
"-q:v", "4", pattern,
]
proc = subprocess.run(cmd, capture_output=True, text=True)
files = sorted(glob.glob(os.path.join(out_dir, "cand_*.jpg")))
if not files:
raise RuntimeError(
"ffmpeg extracted no candidate frames.\n" + proc.stderr[-600:]
)
# fps filter emits frame i (1-based) at ~ (i-1)/eff_fps seconds.
return [((i) / eff_fps, f) for i, f in enumerate(files)]
def delete_video(video_path: str) -> None:
"""Delete the downloaded video (and any sibling temp media) if present."""
for p in (video_path,):
try:
if p and os.path.exists(p):
os.remove(p)
except OSError:
pass
def _clamp(x: float, lo: float, hi: float) -> float:
return max(lo, min(hi, x))
def best_match_segment(quote: str, t_llm: float, segs: list[dict]) -> dict:
"""Find the Whisper segment a step refers to.
Prefer the segment whose text best matches ``quote``; if there's no decent textual
match, fall back to the segment nearest ``t_llm``.
"""
quote = (quote or "").strip().lower()
if quote:
scored = [
(difflib.SequenceMatcher(None, quote, s["text"].lower()).ratio(), s)
for s in segs
]
ratio, best = max(scored, key=lambda t: t[0])
if ratio >= 0.3:
return best
return min(segs, key=lambda s: abs((s["start"] + s["end"]) / 2 - t_llm))
def pick_time(step: dict, segs: list[dict], cand_times: list[float], *,
w_llm: float = 0.4, w_whisper: float = 0.6, lead: float = 1.0,
max_drift: float = 20.0) -> float:
"""Weighted capture time, snapped to the nearest candidate frame.
``t_whisper`` (segment mid + lead) is ground truth for *when* the referenced text is
spoken; ``t_llm`` is the model's guess of *what* matters. We blend them, but if the
LLM time drifts far from the matched segment we clamp it back into that segment so a
hallucinated timestamp can't drag the shot off-topic.
"""
seg = best_match_segment(step.get("quote", ""), float(step.get("t_llm", 0.0)), segs)
t_whisper = (seg["start"] + seg["end"]) / 2 + lead
t_llm = float(step.get("t_llm", t_whisper))
if abs(t_llm - t_whisper) > max_drift:
t_llm = _clamp(t_llm, seg["start"], seg["end"])
total = w_llm + w_whisper or 1.0
t_target = (w_llm * t_llm + w_whisper * t_whisper) / total
return min(cand_times, key=lambda c: abs(c - t_target))
def select_screenshots(steps: list[dict], segs: list[dict],
candidates: list[tuple[float, str]], *,
w_llm: float = 0.4, w_whisper: float = 0.6, lead: float = 1.0,
max_shots: int = 8) -> dict[int, dict]:
"""Return ``{step_index: {"time", "path"}}`` for the steps that get a screenshot.
When there are more steps than ``max_shots``, the highest-``importance`` steps win.
"""
if not candidates:
return {}
cand_times = [t for t, _ in candidates]
time_to_path = dict(candidates)
order = sorted(range(len(steps)), key=lambda i: -float(steps[i].get("importance", 0.5)))
chosen = set(order[:max_shots])
out: dict[int, dict] = {}
for i, step in enumerate(steps):
if i not in chosen:
continue
t = pick_time(step, segs, cand_times, w_llm=w_llm, w_whisper=w_whisper, lead=lead)
out[i] = {"time": t, "path": time_to_path[t]}
return out
|