1"""Frame extraction with ffmpeg: every N seconds, N evenly spaced, exact timestamps, scene changes or keyframes."""
2from __future__ import annotations
3
4import json
5import math
6import os
7import re
8import subprocess
9
10from PIL import Image, ImageDraw, ImageFont
11
12FORMATS = {'jpg': ('jpg', ['-q:v', '3']), 'png': ('png', []), 'webp': ('webp', ['-c:v', 'libwebp', '-quality', '85'])}
13FONT = '/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf'
14
15
16class FrameError(Exception):
17 pass
18
19
20def _run(cmd: list[str], timeout: int = 600) -> subprocess.CompletedProcess:
21 try:
22 return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
23 except subprocess.TimeoutExpired as e:
24 raise FrameError('ffmpeg took too long on this video') from e
25
26
27def probe(path: str) -> dict:
28 p = _run(['ffprobe', '-v', 'error', '-print_format', 'json', '-show_format', '-show_streams', path], 120)
29 try:
30 info = json.loads(p.stdout or '{}')
31 except json.JSONDecodeError:
32 info = {}
33 video = next((s for s in info.get('streams', []) if s.get('codec_type') == 'video'
34 and not (s.get('disposition') or {}).get('attached_pic')), None)
35 if not video:
36 raise FrameError('the file has no video stream')
37 duration = float(info.get('format', {}).get('duration') or video.get('duration') or 0)
38 num, _, den = str(video.get('avg_frame_rate') or '0/1').partition('/')
39 fps = float(num) / float(den or 1) if float(den or 1) else 0.0
40 return {'duration': round(duration, 3), 'width': video.get('width'), 'height': video.get('height'),
41 'fps': round(fps, 3), 'codec': video.get('codec_name'),
42 'formatName': info.get('format', {}).get('format_name')}
43
44
45def parse_time(s: str) -> float:
46 s = s.strip()
47 if not re.fullmatch(r'(\d+:){0,2}\d+(\.\d+)?', s):
48 raise FrameError(f'bad timestamp "{s}" (use seconds like 90.5 or hh:mm:ss)')
49 secs = 0.0
50 for part in s.split(':'):
51 secs = secs * 60 + float(part)
52 return secs
53
54
55def plan_times(mode: str, duration: float, *, interval: float, count: int, timestamps: list[str],
56 max_frames: int) -> list[float]:
57 if duration <= 0 and mode != 'timestamps':
58 raise FrameError('the video duration is unknown; use mode "timestamps"')
59 if mode == 'interval':
60 n = min(max_frames, int(math.floor(duration / interval)) + 1)
61 times = [i * interval for i in range(n)]
62 elif mode == 'count':
63 n = min(max_frames, count)
64 times = [(i + 0.5) * duration / n for i in range(n)]
65 elif mode == 'timestamps':
66 times = [parse_time(t) for t in timestamps][:max_frames]
67 if duration > 0:
68 times = [t for t in times if t < duration]
69 else:
70 raise FrameError(f'unknown mode "{mode}"')
71
72 return [min(t, max(duration - 0.05, 0)) if duration > 0 else t for t in times]
73
74
75def _scale(max_width: int) -> list[str]:
76 return ['-vf', f"scale='min({max_width},iw)':-2"] if max_width else []
77
78
79def grab(path: str, t: float, dest: str, fmt: str, max_width: int) -> bool:
80 _, opts = FORMATS[fmt]
81 p = _run(['ffmpeg', '-v', 'error', '-y', '-ss', f'{t:.3f}', '-i', path, '-frames:v', '1', *_scale(max_width),
82 *opts, '-f', 'image2', '-update', '1', dest], 120)
83 return p.returncode == 0 and os.path.exists(dest) and os.path.getsize(dest) > 0
84
85
86def select_frames(path: str, outdir: str, stem: str, fmt: str, max_width: int, mode: str,
87 threshold: float, max_frames: int) -> list[tuple[float, str]]:
88 """Scene changes or keyframes in one ffmpeg pass. Returns [(time, file)]."""
89 ext, opts = FORMATS[fmt]
90 pattern = os.path.join(outdir, f'{stem}-sel%05d.{ext}')
91
92 filters = ([f"select='eq(n\\,0)+gt(scene\\,{threshold})'"] if mode == 'scenes' else []) + ['showinfo']
93 if max_width:
94 filters.append(f"scale='min({max_width},iw)':-2")
95 pre = ['-skip_frame', 'nokey'] if mode == 'keyframes' else []
96 p = _run(['ffmpeg', '-v', 'info', '-y', *pre, '-i', path, '-vf', ','.join(filters), '-fps_mode', 'vfr',
97 '-frames:v', str(max_frames), *opts, '-f', 'image2', pattern], 1800)
98 if p.returncode != 0:
99 raise FrameError('ffmpeg could not decode the video: ' + (p.stderr.strip().splitlines() or ['?'])[-1][:200])
100 times = [float(m) for m in re.findall(r'pts_time:\s*([0-9.]+)', p.stderr)]
101 files = sorted(f for f in os.listdir(outdir) if f.startswith(f'{stem}-sel'))
102 return [(round(t, 3), os.path.join(outdir, f)) for t, f in zip(times, files)]
103
104
105def timecode(t: float) -> str:
106 h, rem = divmod(t, 3600)
107 m, s = divmod(rem, 60)
108 return f'{int(h):02d}:{int(m):02d}:{s:06.3f}'
109
110
111def contact_sheet(frames: list[tuple[float, str]], dest: str, columns: int = 4, thumb_width: int = 320) -> None:
112 font = ImageFont.truetype(FONT, 16)
113 thumbs = []
114 for t, f in frames:
115 with Image.open(f) as im:
116 im = im.convert('RGB')
117 im.thumbnail((thumb_width, thumb_width * 4))
118 thumbs.append((t, im.copy()))
119 th = max(im.height for _, im in thumbs)
120 rows = math.ceil(len(thumbs) / columns)
121 sheet = Image.new('RGB', (columns * (thumb_width + 8) + 8, rows * (th + 8) + 8), (24, 24, 24))
122 draw = ImageDraw.Draw(sheet)
123 for i, (t, im) in enumerate(thumbs):
124 x, y = 8 + (i % columns) * (thumb_width + 8), 8 + (i // columns) * (th + 8)
125 sheet.paste(im, (x, y))
126 label = timecode(t)[:-4]
127 w = draw.textlength(label, font=font)
128 draw.rectangle([x + 4, y + im.height - 26, x + 12 + w, y + im.height - 4], fill=(0, 0, 0))
129 draw.text((x + 8, y + im.height - 25), label, font=font, fill='white')
130 sheet.save(dest, quality=85)
