169 lines
6.1 KiB
Python
169 lines
6.1 KiB
Python
"""Extract frames from a video into a folder of JPGs.
|
||
|
||
Two engines:
|
||
|
||
* **ffmpeg** (preferred, used when the ``ffmpeg`` binary is on PATH) — one
|
||
subprocess does decode + sampling + optional downscale + JPEG encode, which is
|
||
faster than pulling frames into Python one by one, and unlocks the big win:
|
||
*keyframe-only* extraction (``-skip_frame nokey`` decodes only I-frames, ~10×
|
||
faster than decoding every frame).
|
||
* **OpenCV** fallback (``cv2.VideoCapture``) when ffmpeg is absent.
|
||
|
||
Decoding H.264/HEVC frame-by-frame is inherently the cost; hardware accel doesn't
|
||
help for this (GPU transfer overhead). The only way to be dramatically faster is
|
||
to decode fewer frames — hence the keyframe mode.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import shutil
|
||
import subprocess
|
||
from collections.abc import Callable
|
||
from pathlib import Path
|
||
|
||
import cv2
|
||
|
||
from ..imageio import imwrite_unicode
|
||
|
||
# progress(done_seconds, total_seconds) -> return False to cancel.
|
||
Progress = Callable[[float, float], bool]
|
||
|
||
|
||
def _find_ffmpeg() -> str | None:
|
||
"""ffmpeg on PATH, else the binary bundled with imageio-ffmpeg, else None."""
|
||
exe = shutil.which("ffmpeg")
|
||
if exe:
|
||
return exe
|
||
try:
|
||
import imageio_ffmpeg
|
||
|
||
return imageio_ffmpeg.get_ffmpeg_exe()
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _video_duration_seconds(video_path: str) -> float:
|
||
cap = cv2.VideoCapture(str(video_path))
|
||
try:
|
||
fps = cap.get(cv2.CAP_PROP_FPS) or 0.0
|
||
frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0.0
|
||
return frames / fps if fps > 0 else 0.0
|
||
finally:
|
||
cap.release()
|
||
|
||
|
||
def _quality_to_qscale(jpg_quality: int) -> int:
|
||
"""Map JPEG quality 0..100 to ffmpeg -q:v (2 best .. 31 worst)."""
|
||
q = round(2 + (100 - max(0, min(100, jpg_quality))) / 100 * 29)
|
||
return max(2, min(31, q))
|
||
|
||
|
||
def extract_frames(
|
||
video_path: str,
|
||
out_dir: str,
|
||
step: int = 15,
|
||
keyframes_only: bool = False,
|
||
max_dim: int = 0,
|
||
jpg_quality: int = 92,
|
||
progress: Progress | None = None,
|
||
) -> int:
|
||
"""Save sampled frames of ``video_path`` into ``out_dir`` as JPGs.
|
||
|
||
``keyframes_only`` decodes only keyframes (fast). Otherwise keeps every
|
||
``step``-th frame. ``max_dim`` (>0) caps the longest side. ``progress`` is
|
||
called with (done_seconds, total_seconds); returning ``False`` cancels.
|
||
Returns the number of frames written.
|
||
"""
|
||
out = Path(out_dir)
|
||
out.mkdir(parents=True, exist_ok=True)
|
||
ffmpeg = _find_ffmpeg()
|
||
if ffmpeg:
|
||
return _extract_ffmpeg(ffmpeg, video_path, out, step, keyframes_only, max_dim, jpg_quality, progress)
|
||
return _extract_cv2(video_path, out, step, max_dim, jpg_quality, progress)
|
||
|
||
|
||
# --------------------------------------------------------------------- ffmpeg
|
||
def _build_vf(step: int, keyframes_only: bool, max_dim: int) -> str | None:
|
||
filters: list[str] = []
|
||
if not keyframes_only and step > 1:
|
||
filters.append(f"select=not(mod(n\\,{int(step)}))")
|
||
if max_dim and max_dim > 0:
|
||
# Cap the longest side to max_dim, preserve aspect, never upscale.
|
||
filters.append(f"scale='min({max_dim},iw)':'min({max_dim},ih)':force_original_aspect_ratio=decrease")
|
||
return ",".join(filters) if filters else None
|
||
|
||
|
||
def _extract_ffmpeg(
|
||
ffmpeg: str, video_path: str, out: Path, step: int,
|
||
keyframes_only: bool, max_dim: int, jpg_quality: int, progress: Progress | None,
|
||
) -> int:
|
||
total = _video_duration_seconds(video_path)
|
||
cmd = [ffmpeg, "-hide_banner", "-loglevel", "error", "-nostdin"]
|
||
if keyframes_only:
|
||
cmd += ["-skip_frame", "nokey"] # input option: decode only keyframes
|
||
cmd += ["-i", video_path]
|
||
vf = _build_vf(step, keyframes_only, max_dim)
|
||
if vf:
|
||
cmd += ["-vf", vf]
|
||
cmd += ["-vsync", "0", "-q:v", str(_quality_to_qscale(jpg_quality)),
|
||
"-progress", "pipe:1", str(out / "%06d.jpg")]
|
||
|
||
proc = subprocess.Popen(
|
||
cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
|
||
stdin=subprocess.DEVNULL, text=True, bufsize=1,
|
||
)
|
||
try:
|
||
assert proc.stdout is not None
|
||
for line in proc.stdout:
|
||
if progress is None:
|
||
continue
|
||
line = line.strip()
|
||
if line.startswith("out_time_us=") or line.startswith("out_time_ms="):
|
||
raw = line.split("=", 1)[1]
|
||
try:
|
||
# out_time_us is microseconds; out_time_ms is *also* microseconds
|
||
# in ffmpeg (historical misnomer). Both -> seconds via /1e6.
|
||
done = int(raw) / 1_000_000 if raw.isdigit() else 0.0
|
||
except ValueError:
|
||
done = 0.0
|
||
if progress(done, total) is False:
|
||
proc.terminate()
|
||
break
|
||
finally:
|
||
proc.wait()
|
||
return len(list(out.glob("*.jpg")))
|
||
|
||
|
||
# ---------------------------------------------------------------------- opencv
|
||
def _extract_cv2(
|
||
video_path: str, out: Path, step: int, max_dim: int, jpg_quality: int, progress: Progress | None,
|
||
) -> int:
|
||
cap = cv2.VideoCapture(str(video_path))
|
||
if not cap.isOpened():
|
||
raise RuntimeError(f"Не удалось открыть видео: {video_path}")
|
||
step = max(1, int(step))
|
||
fps = cap.get(cv2.CAP_PROP_FPS) or 0.0
|
||
total = (cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0.0) / fps if fps > 0 else 0.0
|
||
params = [cv2.IMWRITE_JPEG_QUALITY, int(jpg_quality)]
|
||
idx = saved = 0
|
||
try:
|
||
while True:
|
||
if not cap.grab():
|
||
break
|
||
if idx % step == 0:
|
||
ok, frame = cap.retrieve()
|
||
if ok:
|
||
if max_dim and max(frame.shape[:2]) > max_dim:
|
||
s = max_dim / max(frame.shape[:2])
|
||
frame = cv2.resize(frame, None, fx=s, fy=s, interpolation=cv2.INTER_AREA)
|
||
imwrite_unicode(str(out / f"{idx:06d}.jpg"), frame, params)
|
||
saved += 1
|
||
idx += 1
|
||
if progress is not None and idx % 30 == 0:
|
||
done = idx / fps if fps > 0 else 0.0
|
||
if progress(done, total) is False:
|
||
break
|
||
finally:
|
||
cap.release()
|
||
return saved
|