File size: 8,801 Bytes
722bda8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
"""
Frame sampling helpers for the scene pipeline.

For each scene window we want a small set of *representative* frames — enough
for the VLM to read the action, few enough to keep token cost low. The
strategy:

1. Sample N evenly-spaced timestamps inside the window.
2. Extract them with ffmpeg as JPEGs at a reasonable resolution.
3. Drop near-duplicates using a perceptual hash so a static talking head
   doesn't waste 6 frame slots on essentially the same picture.

Each frame is also written to disk under the run's output directory so the
benchmark report can show what the VLM actually saw.
"""

from __future__ import annotations

import os
import shutil
import subprocess
from dataclasses import dataclass
from typing import Sequence

from PIL import Image
import imagehash

from utils import log_message

DEFAULT_FRAMES_PER_WINDOW = 6
DEFAULT_FRAME_HEIGHT = 448
PHASH_DUPLICATE_DISTANCE = 4


@dataclass
class SampledFrame:
    """A frame extracted for a scene window."""

    timestamp_seconds: float
    image_path: str

    def to_dict(self) -> dict:
        return {
            "timestamp_seconds": round(self.timestamp_seconds, 3),
            "image_path": self.image_path,
        }


class FrameExtractionError(RuntimeError):
    """Raised when ffmpeg cannot extract a frame from a video."""


def probe_duration_seconds(video_path: str) -> float:
    """Return the duration of a video in seconds using ffprobe.

    Raises ``FrameExtractionError`` if ffprobe is missing or the video
    cannot be read — never silently returns 0, because a 0 duration would
    cause downstream window construction to produce zero work and the
    operator would have no idea why.
    """
    if shutil.which("ffprobe") is None:
        raise FrameExtractionError(
            "ffprobe is not installed. Install ffmpeg (which includes ffprobe) "
            "before running the scene pipeline."
        )
    if not os.path.exists(video_path):
        raise FrameExtractionError(f"Video file not found: {video_path}")

    result = subprocess.run(
        [
            "ffprobe",
            "-v", "error",
            "-show_entries", "format=duration",
            "-of", "default=noprint_wrappers=1:nokey=1",
            video_path,
        ],
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        raise FrameExtractionError(
            f"ffprobe failed for {video_path}: {result.stderr.strip()}"
        )
    raw = result.stdout.strip()
    if not raw:
        raise FrameExtractionError(f"ffprobe returned no duration for {video_path}")
    try:
        duration = float(raw)
    except ValueError as exc:
        raise FrameExtractionError(
            f"ffprobe returned non-numeric duration {raw!r} for {video_path}"
        ) from exc
    if duration <= 0:
        raise FrameExtractionError(
            f"ffprobe reported non-positive duration {duration} for {video_path}"
        )
    return duration


def sample_window_frames(
    video_path: str,
    window_start: float,
    window_end: float,
    output_dir: str,
    *,
    frames_per_window: int = DEFAULT_FRAMES_PER_WINDOW,
    frame_height: int = DEFAULT_FRAME_HEIGHT,
    dedup_phash: bool = True,
) -> list[SampledFrame]:
    """Extract `frames_per_window` frames from a window and dedup near-duplicates.

    Args:
        video_path: Source mp4 path.
        window_start: Window start in seconds.
        window_end: Window end in seconds (exclusive).
        output_dir: Directory to write extracted JPEGs into. Created if missing.
        frames_per_window: Target frame count before deduplication.
        frame_height: Resize height in pixels. Width is auto (preserves aspect).
        dedup_phash: Drop frames whose perceptual hash is within
            PHASH_DUPLICATE_DISTANCE bits of an earlier kept frame.

    Returns:
        Frames kept after deduplication, ordered by timestamp.
    """
    if shutil.which("ffmpeg") is None:
        raise FrameExtractionError(
            "ffmpeg is not installed. Install ffmpeg before running the scene pipeline."
        )
    if window_end <= window_start:
        return []
    if frames_per_window <= 0:
        return []

    os.makedirs(output_dir, exist_ok=True)
    timestamps = _evenly_spaced_timestamps(window_start, window_end, frames_per_window)

    extracted: list[SampledFrame] = []
    for slot_index, timestamp in enumerate(timestamps):
        image_path = os.path.join(output_dir, f"frame-{slot_index:02d}.jpg")
        try:
            _extract_single_frame(
                video_path=video_path,
                timestamp_seconds=timestamp,
                output_path=image_path,
                frame_height=frame_height,
            )
        except FrameExtractionError as exc:
            log_message(
                f"Skipping frame at {timestamp:.2f}s of {video_path}: {exc}",
                "WARN",
            )
            continue
        extracted.append(SampledFrame(timestamp_seconds=timestamp, image_path=image_path))

    if not dedup_phash or len(extracted) <= 1:
        return extracted

    return _dedup_by_phash(extracted)


def _evenly_spaced_timestamps(start: float, end: float, count: int) -> list[float]:
    """Return N timestamps evenly spaced *inside* (start, end).

    Avoids the exact endpoints — boundary frames are often cuts or fades
    that aren't representative of the window's content.
    """
    if count == 1:
        return [start + (end - start) / 2.0]
    step = (end - start) / (count + 1)
    return [start + step * (i + 1) for i in range(count)]


def _extract_single_frame(
    *,
    video_path: str,
    timestamp_seconds: float,
    output_path: str,
    frame_height: int,
) -> None:
    """Run ffmpeg to grab one frame. Fast-seek before -i for speed."""
    cmd = [
        "ffmpeg",
        "-y",
        "-ss", f"{timestamp_seconds:.3f}",
        "-i", video_path,
        "-frames:v", "1",
        "-vf", f"scale=-2:{frame_height}",
        "-q:v", "3",
        "-loglevel", "error",
        output_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, check=False)
    if result.returncode != 0 or not os.path.exists(output_path):
        raise FrameExtractionError(
            f"ffmpeg failed at {timestamp_seconds}s: {result.stderr.strip()}"
        )


def _dedup_by_phash(frames: Sequence[SampledFrame]) -> list[SampledFrame]:
    """Drop near-duplicate frames using perceptual hash distance."""
    kept: list[SampledFrame] = []
    kept_hashes: list[imagehash.ImageHash] = []
    for frame in frames:
        try:
            with Image.open(frame.image_path) as img:
                phash = imagehash.phash(img)
        except (OSError, ValueError) as exc:
            log_message(
                f"phash failed for {frame.image_path}: {exc}; keeping frame anyway",
                "WARN",
            )
            kept.append(frame)
            continue

        is_duplicate = any(
            (phash - prior) <= PHASH_DUPLICATE_DISTANCE for prior in kept_hashes
        )
        if is_duplicate:
            try:
                os.remove(frame.image_path)
            except OSError:
                pass
            continue
        kept.append(frame)
        kept_hashes.append(phash)
    return kept


def extract_audio_clip(
    video_path: str,
    window_start: float,
    window_end: float,
    output_path: str,
    *,
    sample_rate: int = 48000,
) -> str:
    """Extract a window's audio as a mono WAV file for CLAP.

    CLAP expects 48kHz mono. We resample at extraction time so the audio
    tagger doesn't need to do it per-call.

    Returns the output path on success. Raises ``FrameExtractionError``
    on failure rather than returning silently — a missing audio clip
    should be visible in the report, not papered over.
    """
    if shutil.which("ffmpeg") is None:
        raise FrameExtractionError("ffmpeg is not installed.")
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    duration = max(0.0, window_end - window_start)
    if duration <= 0:
        raise FrameExtractionError(
            f"Non-positive window duration ({duration}s) for audio extraction"
        )

    cmd = [
        "ffmpeg",
        "-y",
        "-ss", f"{window_start:.3f}",
        "-t", f"{duration:.3f}",
        "-i", video_path,
        "-vn",
        "-ac", "1",
        "-ar", str(sample_rate),
        "-loglevel", "error",
        output_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, check=False)
    if result.returncode != 0 or not os.path.exists(output_path):
        raise FrameExtractionError(
            f"ffmpeg audio extract failed at {window_start}s: {result.stderr.strip()}"
        )
    return output_path