src/extract_frames.py

"""Extract still frames from a video at given timestamps via ffmpeg."""

from __future__ import annotations

import json
import os
import subprocess
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

DEFAULT_WORKERS = 6


def _workers() -> int:
    """Parallel ffmpeg processes. Each is CPU-hungry, so keep this modest."""
    raw = (os.getenv("FFMPEG_WORKERS") or "").strip()
    return int(raw) if raw.isdigit() and int(raw) > 0 else DEFAULT_WORKERS


def probe_duration_sec(video_path: Path | str) -> float:
    """Return media duration in seconds via ffprobe."""
    video = Path(video_path).expanduser().resolve()
    if not video.is_file():
        raise FileNotFoundError(f"Video not found: {video}")
    cmd = [
        "ffprobe",
        "-v",
        "error",
        "-show_entries",
        "format=duration",
        "-of",
        "json",
        str(video),
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0:
        err = (proc.stderr or proc.stdout or "").strip()[-400:]
        raise RuntimeError(f"ffprobe failed: {err}")
    data = json.loads(proc.stdout or "{}")
    duration = float((data.get("format") or {}).get("duration") or 0)
    if duration <= 0:
        raise RuntimeError(f"Could not read duration for {video}")
    return duration


def clamp_timestamp(ts: float, duration_sec: float, *, margin: float = 0.15) -> float:
    """Keep seek inside [0, duration - margin] so ffmpeg can decode a frame."""
    if duration_sec <= 0:
        return max(0.0, float(ts))
    end = max(0.0, duration_sec - margin)
    return min(max(0.0, float(ts)), end)


def extract_frames(
    video_path: Path | str,
    timestamps_sec: list[float],
    out_dir: Path | str,
    *,
    prefix: str = "frame",
    duration_sec: float | None = None,
) -> list[Path]:
    """Write one JPEG per timestamp into out_dir. Returns paths in order."""
    video = Path(video_path).expanduser().resolve()
    if not video.is_file():
        raise FileNotFoundError(f"Video not found: {video}")
    dest = Path(out_dir)
    dest.mkdir(parents=True, exist_ok=True)

    duration = duration_sec if duration_sec is not None else probe_duration_sec(video)

    def grab(item: tuple[int, float]) -> Path:
        i, ts = item
        t = clamp_timestamp(ts, duration)
        out = dest / f"{prefix}_{i:02d}.jpg"
        tail = ["-frames:v", "1", "-q:v", "2", "-an", str(out)]
        # Input seek (-ss before -i) jumps straight to the nearest keyframe instead of
        # decoding from t=0; falls back to output seek, which is reliable past the last one.
        for cmd in (
            ["ffmpeg", "-nostdin", "-y", "-ss", f"{t:.3f}", "-i", str(video), *tail],
            ["ffmpeg", "-nostdin", "-y", "-i", str(video), "-ss", f"{t:.3f}", *tail],
        ):
            proc = subprocess.run(cmd, capture_output=True, text=True)
            if proc.returncode == 0 and out.is_file() and out.stat().st_size > 0:
                return out
        err = (proc.stderr or proc.stdout or "").strip()[-800:]
        raise RuntimeError(f"ffmpeg failed at t={t:.3f}s (duration={duration:.3f}s): {err}")

    items = list(enumerate(timestamps_sec))
    if len(items) < 2:
        return [grab(it) for it in items]
    with ThreadPoolExecutor(max_workers=min(_workers(), len(items))) as pool:
        return list(pool.map(grab, items))  # ordered; first exception propagates